QLoRA Advanced: Multi-Adapter Training - Ministral-3B-Reasoning¶

# Environment Setup
import os
from dotenv import load_dotenv
load_dotenv()

# Force text-based progress instead of HTML widgets
os.environ["TQDM_NOTEBOOK"] = "false"

# CRITICAL: Import unsloth FIRST for proper TRL patching
import unsloth
from unsloth import FastLanguageModel, is_bf16_supported

import torch
import gc

# Environment summary
gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
print(f"Environment: unsloth {unsloth.__version__}, PyTorch {torch.__version__}, {gpu}")
print(f"HF_TOKEN loaded: {'Yes' if os.environ.get('HF_TOKEN') else 'No'}")

# Helper Functions

def cleanup_memory():
    """Force garbage collection and clear CUDA cache"""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

def count_parameters(model):
    """Count trainable vs total parameters"""
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    return {"trainable": trainable, "total": total, "pct": 100 * trainable / total}

print("Helper functions defined.")

# Task-Specific Datasets (5 samples each)
from datasets import Dataset

# Task 1: Technical Q&A - precise, factual responses
technical_data = [
    {
        "instruction": "What is a hash table?",
        "thinking": "The user wants a technical explanation. What are the key aspects? It's a data structure... uses hashing for O(1) lookups... handles collisions. Keep it precise but accessible.",
        "response": "A hash table is a data structure that maps keys to values using a hash function. It provides O(1) average-case time complexity for insertions, deletions, and lookups by computing an index from the key."
    },
    {
        "instruction": "Explain TCP vs UDP.",
        "thinking": "Compare two network protocols. TCP is reliable, connection-oriented. UDP is fast, connectionless. What are the key differences to highlight?",
        "response": "TCP provides reliable, ordered delivery with connection establishment and error checking. UDP is connectionless and faster but unreliable. TCP is used for web traffic; UDP for streaming and gaming."
    },
    {
        "instruction": "What is Big O notation?",
        "thinking": "Fundamental CS concept. It describes algorithmic complexity. How do I explain it simply? Focus on growth rate and worst-case analysis.",
        "response": "Big O notation describes the upper bound of an algorithm's time or space complexity as input size grows. O(n) means linear growth, O(n²) means quadratic, and O(1) means constant time."
    },
    {
        "instruction": "What is recursion?",
        "thinking": "Classic programming concept. A function calling itself with a base case. Should I mention stack frames? Keep it conceptual first.",
        "response": "Recursion is when a function calls itself to solve smaller instances of the same problem. It requires a base case to terminate and recursive cases that reduce toward the base case."
    },
    {
        "instruction": "Explain database indexing.",
        "thinking": "Database performance topic. Indexes speed up queries... B-trees are common. Trade-off with insert performance. Focus on the core concept.",
        "response": "Database indexing creates data structures that speed up query operations by allowing the database to find rows without scanning entire tables. Indexes trade storage space and insert speed for faster reads."
    },
]

# Task 2: Creative Writing - imaginative, expressive responses
creative_data = [
    {
        "instruction": "Write a haiku about coding.",
        "thinking": "Haiku: 5-7-5 syllables. Theme is coding. What imagery captures the experience? Bugs, late nights, the flow state... Let me craft something evocative.",
        "response": "Lines of logic flow\nBugs dance in the midnight glow\nFixed—the dawn says so"
    },
    {
        "instruction": "Describe a sunset poetically.",
        "thinking": "Poetry mode. What makes sunsets beautiful? Colors, transition, emotion. Use metaphor and sensory language.",
        "response": "The sun dips low, painting the sky in watercolors of amber and rose. Day surrenders to twilight's embrace as the horizon glows like embers of a dying fire, beautiful in its farewell."
    },
    {
        "instruction": "Create a metaphor for learning.",
        "thinking": "What is learning like? A journey? Growing? Building? Something that captures the gradual, sometimes difficult process...",
        "response": "Learning is climbing a mountain in fog—each step reveals new terrain, the path unclear until you've walked it. The summit exists, but the journey shapes you more than arrival."
    },
    {
        "instruction": "Write a short opening line for a mystery novel.",
        "thinking": "Mystery opening needs intrigue. Something that raises questions. A hook that makes readers want more. Tension from the first sentence.",
        "response": "The letter arrived three years after her death, written in handwriting that was unmistakably hers."
    },
    {
        "instruction": "Describe rain in an unusual way.",
        "thinking": "Avoid clichés. What's a fresh perspective on rain? Sound, feeling, unusual comparisons... Make it memorable.",
        "response": "The sky plays percussion on the rooftops—a million tiny drummers performing chaos into rhythm, washing the world in liquid applause."
    },
]

# Task 3: Code Explanation - code-focused analysis
code_data = [
    {
        "instruction": "Explain: def fib(n): return n if n<2 else fib(n-1)+fib(n-2)",
        "thinking": "Classic recursive Fibonacci. Base case n<2, recursive case adds previous two. What should I highlight? The elegance but also the O(2^n) complexity.",
        "response": "This is a recursive Fibonacci implementation. Base case: returns n for 0 or 1. Recursive case: sums the two preceding values. Elegant but O(2^n) time—use memoization for efficiency."
    },
    {
        "instruction": "What does [x*2 for x in range(5)] do?",
        "thinking": "Python list comprehension. Creates a list of doubled values. range(5) gives 0-4. Result is [0,2,4,6,8]. Explain the syntax pattern.",
        "response": "This list comprehension creates [0, 2, 4, 6, 8]. It iterates x through range(5) (0-4) and multiplies each by 2. Syntax: [expression for item in iterable]."
    },
    {
        "instruction": "Explain Python decorators.",
        "thinking": "Decorators wrap functions. They're higher-order functions. @syntax is syntactic sugar. Use cases: logging, timing, auth. Start with the concept.",
        "response": "Decorators are functions that modify other functions. The @decorator syntax wraps a function, adding behavior before/after execution. Common uses: logging, authentication, caching, timing."
    },
    {
        "instruction": "What is async/await in Python?",
        "thinking": "Asynchronous programming. Coroutines, event loop, non-blocking I/O. How do I explain this concisely? Focus on the why and basic usage.",
        "response": "async/await enables non-blocking I/O. 'async def' creates coroutines; 'await' pauses execution until a task completes. The event loop manages concurrent operations without threads."
    },
    {
        "instruction": "Explain: lambda x,y: x if x>y else y",
        "thinking": "Lambda function returning the max of two values. Anonymous function, single expression. Equivalent to max() for two values.",
        "response": "This lambda returns the larger of x and y using a ternary expression. It's equivalent to max(x,y). Lambda creates anonymous, single-expression functions."
    },
]

TASK_DATASETS = {
    "technical": technical_data,
    "creative": creative_data,
    "code": code_data,
}

print(f"Task datasets prepared:")
for name, data in TASK_DATASETS.items():
    print(f"  - {name}: {len(data)} samples")

# Adapter Switching Demonstration
from peft import PeftModel

print("="*60)
print("Adapter Switching Demonstration")
print("="*60)
print("\nSame prompt with different adapters:")
print("Prompt: 'Explain what a variable is.'")
print("="*60)

TEST_PROMPT = "Explain what a variable is."
results = {}

# Load base model once
cleanup_memory()
print("\nLoading base model...")
base_model, tokenizer = FastLanguageModel.from_pretrained(
    MODEL_NAME,
    max_seq_length=512,
    load_in_4bit=True,
    dtype=None,
)

# Use text tokenizer directly for text-only inference (Pixtral uses a processor)
text_tokenizer = tokenizer.tokenizer if hasattr(tokenizer, 'tokenizer') else tokenizer

for task_name, adapter_path in adapter_paths.items():
    print(f"\n--- Loading {task_name} adapter ---")
    
    # Load adapter on base model
    adapted_model = PeftModel.from_pretrained(base_model, adapter_path)
    FastLanguageModel.for_inference(adapted_model)
    
    # Generate response - using Ministral multimodal format
    messages = [{"role": "user", "content": [{"type": "text", "text": TEST_PROMPT}]}]
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = text_tokenizer(prompt, return_tensors="pt").to("cuda")
    
    with torch.no_grad():
        outputs = adapted_model.generate(
            **inputs,
            max_new_tokens=150,
            temperature=0.6,
            do_sample=True,
            pad_token_id=text_tokenizer.pad_token_id,
        )
    
    response = text_tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    results[task_name] = response
    
    print(f"Response preview: {response[:200]}...")
    
    # Unload adapter
    del adapted_model
    cleanup_memory()

# Cleanup base model
del base_model, tokenizer
cleanup_memory()

print(f"\n{'='*60}")
print("Demonstration complete!")
print(f"{'='*60}")

# Compare Responses Side by Side
print("="*60)
print("Response Comparison: Same Prompt, Different Adapters")
print("="*60)
print(f"\nPrompt: '{TEST_PROMPT}'\n")

for task_name, response in results.items():
    print(f"\n[{task_name.upper()} ADAPTER]")
    print("-" * 40)
    # Show full response
    print(response)
    print()

# Shutdown kernel to release all GPU memory
import IPython
print("Shutting down kernel to release GPU memory...")
app = IPython.Application.instance()
app.kernel.do_shutdown(restart=False)