SFT Training Test: Ministral (Text-Only)¶

# 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

# 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'}")

# Load Ministral-3B with 4-bit quantization (using FastLanguageModel for text-only)
MODEL_NAME = "unsloth/Ministral-3-3B-Reasoning-2512"
print(f"\nLoading {MODEL_NAME.split('/')[-1]} with FastLanguageModel (text-only mode)...")

model, tokenizer = FastLanguageModel.from_pretrained(
    MODEL_NAME,
    max_seq_length=512,
    load_in_4bit=True,
    dtype=None,  # Auto-detect
)
print(f"Model loaded: {type(model).__name__}")

# Apply LoRA adapters
model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    lora_alpha=16,
    lora_dropout=0,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=42,
)

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"LoRA applied: {trainable:,} trainable / {total:,} total ({100*trainable/total:.2f}%)")

# Create minimal synthetic instruction dataset (5 samples)
# Using Ministral's multimodal chat format for text-only content
from datasets import Dataset

# Synthetic instruction-response pairs for testing
synthetic_data = [
    {"instruction": "What is machine learning?",
     "response": "Machine learning is a subset of artificial intelligence where computers learn patterns from data."},
    {"instruction": "Explain Python in one sentence.",
     "response": "Python is a high-level programming language known for its readability and versatility."},
    {"instruction": "What is a neural network?",
     "response": "A neural network is a computational model inspired by biological neurons that processes information."},
    {"instruction": "Define supervised learning.",
     "response": "Supervised learning is training a model on labeled data to predict outcomes for new inputs."},
    {"instruction": "What is gradient descent?",
     "response": "Gradient descent is an optimization algorithm that minimizes loss by iteratively adjusting parameters."},
]

# Format as chat conversations using Ministral's multimodal format
def format_conversation(sample):
    # Ministral uses multimodal format even for text: [{"type": "text", "text": "..."}]
    messages = [
        {"role": "user", "content": [{"type": "text", "text": sample["instruction"]}]},
        {"role": "assistant", "content": [{"type": "text", "text": sample["response"]}]}
    ]
    return {"text": tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)}

dataset = Dataset.from_list(synthetic_data)
dataset = dataset.map(format_conversation, remove_columns=["instruction", "response"])
print(f"Dataset created: {len(dataset)} samples")
print(f"Sample: {dataset[0]['text'][:150]}...")

# Post-training inference test
FastLanguageModel.for_inference(model)

# Use Ministral's multimodal format for text
messages = [{"role": "user", "content": [{"type": "text", "text": "What is deep learning?"}]}]
input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(None, input_text, add_special_tokens=False, return_tensors="pt").to("cuda")

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=64,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tokenizer.pad_token_id,
    )

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Clean up BPE artifacts from Ministral tokenizer (Ġ=space, Ċ=newline)
response = response.replace('Ġ', ' ').replace('Ċ', '\n').strip()

# Clear success/failure banner
print("=" * 60)
if SFT_TEXT_SUPPORTED:
    print("SFT Training: SUPPORTED for Ministral (Text-Only)")
    print("Model: FastLanguageModel + Ministral-3-3B-Reasoning-2512")
else:
    print("SFT Training: NOT SUPPORTED for Ministral (Text-Only)")
    print("Reason: See error above")
print("=" * 60)
print(f"Sample generation:\n{response[-200:]}")

# 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)