# ============================================
# 🎖️ CYBERRANGER V22 - THE REFINED MIND
# Complete Training & Conversion Pipeline
# ============================================
# Run this ENTIRE notebook in Google Colab
# Output: GGUF file ready for Ollama
# ============================================
# V22 = Precision + Identity + Quality + Fun
# ============================================

# ============================================
# CELL 1: Install Dependencies
# ============================================
!pip install -q transformers datasets accelerate peft bitsandbytes trl
!pip install -q huggingface_hub
!pip install -q sentencepiece

# Clone llama.cpp for GGUF conversion
!git clone --depth 1 https://github.com/ggerganov/llama.cpp
!pip install -q -r llama.cpp/requirements.txt

print("✅ Dependencies installed!")

# ============================================
# CELL 2: Upload Training Data
# ============================================
# Upload qbrain_training_v22_refined.json from your computer
from google.colab import files
print("📤 Upload qbrain_training_v22_refined.json")
uploaded = files.upload()

# ============================================
# CELL 3: Load and Prepare Data
# ============================================
import json
from datasets import Dataset

# Load training data
with open('qbrain_training_v22_refined.json', 'r') as f:
    training_data = json.load(f)

print(f"📊 Loaded {len(training_data)} training examples")

# Convert to HuggingFace Dataset format
def format_for_training(examples):
    """Format as instruction-output pairs."""
    texts = []
    for ex in examples:
        text = f"<|im_start|>user
{ex['instruction']}<|im_end|>
<|im_start|>assistant
{ex['output']}<|im_end|>"
        texts.append(text)
    return texts

formatted_texts = format_for_training(training_data)
dataset = Dataset.from_dict({"text": formatted_texts})

print(f"✅ Dataset prepared: {len(dataset)} examples")
print(f"📝 Sample:
{dataset[0]['text'][:500]}...")

# ============================================
# CELL 4: Load Base Model with QLoRA
# ============================================
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# Model configuration
MODEL_NAME = "HuggingFaceTB/SmolLM2-1.7B-Instruct"

# QLoRA config (4-bit quantization for training)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

print(f"🔄 Loading {MODEL_NAME}...")

# Load model
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

print("✅ Base model loaded!")

# ============================================
# CELL 5: Configure LoRA
# ============================================
# Prepare model for training
model = prepare_model_for_kbit_training(model)

# LoRA configuration - targeting key layers
lora_config = LoraConfig(
    r=16,                          # Rank
    lora_alpha=16,                 # Alpha
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",  # Attention
        "gate_proj", "up_proj", "down_proj",     # MLP
        "embed_tokens", "lm_head"                # Embeddings
    ],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, lora_config)

# Print trainable parameters
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in model.parameters())
print(f"🧠 Trainable: {trainable_params:,} / {total_params:,} ({100 * trainable_params / total_params:.2f}%)")

# ============================================
# CELL 6: Training Configuration
# ============================================
from trl import SFTTrainer, SFTConfig

# Training arguments
training_args = SFTConfig(
    output_dir="./cyberranger_v22",
    num_train_epochs=5,             # Higher epochs for V22 to ensure precision anchoring
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=1e-4,             # Lower LR for better refinement
    weight_decay=0.01,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    logging_steps=10,
    save_steps=100,
    save_total_limit=2,
    fp16=True,
    optim="paged_adamw_8bit",
    max_seq_length=512,
    dataset_text_field="text",
    packing=False,
)

# Create trainer
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
)

print("✅ Trainer configured!")
print(f"📊 Training for {training_args.num_train_epochs} epochs")

# ============================================
# CELL 7: TRAIN! 🚀
# ============================================
print("🚀 Starting V22 REFINED MIND training...")
print("=" * 50)

trainer.train()

print("=" * 50)
print("✅ Training complete!")

# Save the adapter
trainer.save_model("./CyberRanger_V22_Adapters")
tokenizer.save_pretrained("./CyberRanger_V22_Adapters")

print("💾 Adapter saved to ./CyberRanger_V22_Adapters")

# ============================================
# CELL 8: Merge Adapter with Base Model
# ============================================
print("🔄 Merging adapter with base model...")

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

# Reload base model (full precision for merging)
base_model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True,
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)

# Load and merge adapter
model = PeftModel.from_pretrained(base_model, "./CyberRanger_V22_Adapters")
merged_model = model.merge_and_unload()

# Save merged model
merged_model.save_pretrained("./merged_v22", safe_serialization=True)
tokenizer.save_pretrained("./merged_v22")

print("✅ Merged model saved to ./merged_v22")

# ============================================
# CELL 9: Convert to GGUF
# ============================================
print("🔄 Converting to GGUF format...")

!python llama.cpp/convert_hf_to_gguf.py ./merged_v22 --outfile cyberranger-v22-f16.gguf --outtype f16

print("✅ GGUF created: cyberranger-v22-f16.gguf")

# ============================================
# CELL 10: Quantize
# ============================================
print("🔄 Quantizing to Q4_K_M...")
!cd llama.cpp && make llama-quantize
!./llama.cpp/llama-quantize cyberranger-v22-f16.gguf cyberranger-v22-q4.gguf q4_k_m

# ============================================
# CELL 11: Download Files
# ============================================
from google.colab import files
files.download('cyberranger-v22-q4.gguf')

print("✅ Download complete!")

# ============================================
# CELL 12: Summary
# ============================================
print("""
============================================
🎖️ V22 REFINED MIND TRAINING COMPLETE!
============================================
REFINEMENTS INCLUDED:
1. Precision Spelling (Strawberry = 3 r's)
2. Creator Distinction (Keane vs Beckham)
3. Fixed Keane Ratios (73.60/27.19/7.57)
4. Curated Joke Library
5. Clean Output (Reduced Emoji Density)

NEXT STEPS:
1. Move GGUF to qbrain folder
2. Create Ollama model with Modelfile.v22-refined
3. Run: ollama run rangerbot:v22

Rangers lead the way! 🎖️🔥🚀
============================================
""")