import json

# Your training examples
examples = [
    {
        "messages": [
            {"role": "system", "content": "You are a pet breed expert."},
            {"role": "user", "content": "What breed is a small white dog with curly hair?"},
            {"role": "assistant", "content": "That's a Bichon Frise! 🐩 They're 9.5-11.5 inches tall, weigh 12-18 lbs, and have a hypoallergenic coat that needs grooming every 4-6 weeks. Wonderful family dogs!"}
        ]
    },
    {
        "messages": [
            {"role": "system", "content": "You are a pet breed expert."},
            {"role": "user", "content": "I need a large, calm dog that's good with kids."},
            {"role": "assistant", "content": "A Golden Retriever or Bernese Mountain Dog would be perfect! 🐕 Both are gentle giants — calm temperament, patient with children, and highly trainable. Goldens are more active; Bernese are couch potatoes."}
        ]
    }
    # ... add 100-500+ examples for best results
]

# Save as JSONL
with open("training_data.jsonl", "w") as f:
    for example in examples:
        f.write(json.dumps(example) + "\n")

print(f"✅ Saved {len(examples)} examples to training_data.jsonl")

# Create a bucket
gsutil mb gs://your-gemma-finetune-bucket

# Upload your dataset
gsutil cp training_data.jsonl gs://your-gemma-finetune-bucket/data/

# Upload validation set (optional but recommended)
gsutil cp validation_data.jsonl gs://your-gemma-finetune-bucket/data/

# Create a virtual environment
python -m venv gemma-env
source gemma-env/bin/activate

# Install the magic stack
pip install torch>=2.2.0
pip install transformers>=4.40.0
pip install trl>=0.8.0
pip install peft>=0.10.0
pip install datasets
pip install accelerate
pip install bitsandbytes  # For QLoRA (4-bit quantization)
pip install google-cloud-storage

# Install the gcloud CLI if you haven't
curl https://sdk.cloud.google.com | bash

# Authenticate
gcloud auth login
gcloud config set project YOUR_PROJECT_ID

# Enable required APIs
gcloud services enable run.googleapis.com
gcloud services enable artifactregistry.googleapis.com
gcloud services enable cloudbuild.googleapis.com

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    BitsAndBytesConfig,
)
from peft import LoraConfig, get_peft_model
from trl import SFTTrainer
from datasets import load_dataset

# ============================================
# 🔧 CONFIGURATION — Tweak these!
# ============================================

MODEL_ID = "google/gemma-4-9b-it"  # Base model
DATASET_PATH = "training_data.jsonl"
OUTPUT_DIR = "./gemma-4-finetuned"

# LoRA config — the secret sauce 🧪
LORA_CONFIG = LoraConfig(
    r=16,                    # Rank (higher = more capacity, more VRAM)
    lora_alpha=32,           # Scaling factor
    lora_dropout=0.05,       # Regularization
    target_modules=[         # Which layers to adapt
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ],
    bias="none",
    task_type="CAUSAL_LM",
)

# Training hyperparameters
TRAINING_ARGS = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=3,              # 3 epochs is usually the sweet spot
    per_device_train_batch_size=2,   # Adjust based on VRAM
    gradient_accumulation_steps=8,   # Effective batch size = 2 * 8 = 16
    learning_rate=2e-4,              # LoRA likes higher LR than full FT
    warmup_ratio=0.1,
    weight_decay=0.01,
    logging_steps=10,
    save_strategy="epoch",
    evaluation_strategy="epoch",
    fp16=True,                       # Mixed precision for speed
    optim="paged_adamw_8bit",        # Memory-efficient optimizer
    gradient_checkpointing=True,     # Save VRAM at cost of speed
    max_grad_norm=1.0,
    report_to="none",                # Change to "wandb" if you use it
)

# ============================================
# 🚀 TRAINING CODE
# ============================================

print("📦 Loading model...")
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=quantization_config,
    device_map="auto",
    torch_dtype=torch.float16,
    attn_impl="flash_attention_2",  # Use Flash Attention if available
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

print("🔧 Applying LoRA...")
model = get_peft_model(model, LORA_CONFIG)
model.print_trainable_parameters()
# Output: trainable params: 41,943,040 || all params: 9,284,536,320 || trainable%: 0.45%

print("📊 Loading dataset...")
dataset = load_dataset("json", data_files=DATASET_PATH, split="train")

print("🚀 Starting training...")
trainer = SFTTrainer(
    model=model,
    args=TRAINING_ARGS,
    train_dataset=dataset,
    tokenizer=tokenizer,
    packing=True,              # Pack short examples together for efficiency
    max_seq_length=2048,       # Max tokens per example
)

trainer.train()

print("💾 Saving adapter...")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

print("✅ Done! Adapter saved to", OUTPUT_DIR)

┌─────────────────────────────────────────────────────┐
│  LoRA Rank (r)                                       │
│  ─────────────                                       │
│  r=8   → Faster, less VRAM, might underfit          │
│  r=16  → Sweet spot for most tasks ⭐               │
│  r=32  → More capacity, needs more data             │
│  r=64  → Diminishing returns, use full FT instead   │
│                                                      │
│  Target Modules                                      │
│  ──────────────                                      │
│  q_proj, v_proj only → Minimum adaptation            │
│  All attention layers → Recommended ⭐               │
│  + MLP layers → Maximum adaptation (more VRAM)      │
└─────────────────────────────────────────────────────┘

FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04

# Install Python
RUN apt-get update && apt-get install -y python3 python3-pip python3-venv

# Set working directory
WORKDIR /app

# Copy requirements and install
COPY requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt

# Copy training code
COPY train.py .
COPY data/ ./data/

# Run training
CMD ["python3", "train.py"]

torch>=2.2.0
transformers>=4.40.0
trl>=0.8.0
peft>=0.10.0
datasets
accelerate
bitsandbytes
google-cloud-storage
flash-attn>=2.5.0

# Build the container
gcloud builds submit --tag gcr.io/YOUR_PROJECT_ID/gemma-finetune

# Create a Cloud Run Job with GPU
gcloud run jobs create gemma-finetune-job \
    --image gcr.io/YOUR_PROJECT_ID/gemma-finetune \
    --region us-central1 \
    --gpu 1 \
    --gpu-type nvidia-l4 \
    --memory 32Gi \
    --cpu 8 \
    --task-timeout 14400 \
    --max-retries 0 \
    --set-env-vars "MODEL_ID=google/gemma-4-9b-it" \
    --service-account YOUR_SERVICE_ACCOUNT@YOUR_PROJECT.iam.gserviceaccount.com

# 🚀 Launch the job!
gcloud run jobs execute gemma-finetune-job --region us-central1

# Watch the logs in real-time
gcloud run jobs executions list --job gemma-finetune-job --region us-central1

# Get the latest execution
EXECUTION=$(gcloud run jobs executions list \
    --job gemma-finetune-job \
    --region us-central1 \
    --format="value(name)" \
    --limit=1)

# Stream logs
gcloud beta run jobs executions logs read $EXECUTION --region us-central1

📦 Loading model...
🔧 Applying LoRA...
trainable params: 41,943,040 || all params: 9,284,536,320 || trainable%: 0.45%
📊 Loading dataset...
🚀 Starting training...

{'loss': 2.3456, 'learning_rate': 0.0002, 'epoch': 0.33}
{'loss': 1.8234, 'learning_rate': 0.00018, 'epoch': 0.67}
{'loss': 1.4567, 'learning_rate': 0.00016, 'epoch': 1.0}
...

✅ Done! Adapter saved to ./gemma-4-finetuned

Loss
 │
2.5┤ ●
   │  ●
2.0┤    ●
   │      ●
1.5┤        ●  ●
   │             ●  ●
1.0┤                   ●  ●  ●    ← Converging nicely! ✅
   │
0.5┤
   └──────────────────────────────
   0    0.5    1.0    1.5    2.0
                Epoch

🚨 Warning signs:
   • Loss stays flat → Learning rate too low
   • Loss explodes → Learning rate too high
   • Train ↓ but val ↑ → Overfitting!

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

# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-9b-it",
    torch_dtype=torch.float16,
    device_map="auto",
)

# Load your fine-tuned adapter
model = PeftModel.from_pretrained(base_model, "./gemma-4-finetuned")
tokenizer = AutoTokenizer.from_pretrained("./gemma-4-finetuned")

# Test it!
def ask(question, system_prompt="You are a pet breed expert."):
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": question},
    ]
    input_text = tokenizer.apply_chat_template(messages, tokenize=False)
    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=200,
            temperature=0.7,
            top_p=0.9,
        )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("model\n")[-1].strip()

# 🧪 Test questions
questions = [
    "What breed is a small white dog with curly hair?",
    "I need a large, calm dog that's good with kids.",
    "Which dog breed is best for apartments?",
    "What's the difference between a Husky and a Malamute?",
]

for q in questions:
    print(f"\n❓ {q}")
    print(f"🐕 {ask(q)}")
    print("-" * 60)

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

print("📦 Loading base model...")
base_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-9b-it",
    torch_dtype=torch.float16,
    device_map="cpu",
)

print("🔗 Merging LoRA adapter...")
model = PeftModel.from_pretrained(base_model, "./gemma-4-finetuned")
merged_model = model.merge_and_unload()  # Merge adapter into base weights

print("💾 Saving merged model...")
merged_model.save_pretrained("./gemma-4-merged")
tokenizer = AutoTokenizer.from_pretrained("./gemma-4-finetuned")
tokenizer.save_pretrained("./gemma-4-merged")

print("📤 Uploading to GCS...")
import subprocess
subprocess.run([
    "gsutil", "-m", "cp", "-r",
    "./gemma-4-merged",
    "gs://your-gemma-finetune-bucket/models/"
])

print("✅ Merged model uploaded!")

# Deploy a vLLM endpoint on Cloud Run
gcloud run deploy gemma-4-api \
    --image vllm/vllm-openai:latest \
    --region us-central1 \
    --gpu 1 \
    --gpu-type nvidia-l4 \
    --memory 32Gi \
    --cpu 8 \
    --allow-unauthenticated \
    --set-env-vars "MODEL=gs://your-gemma-finetune-bucket/models/gemma-4-merged"

curl -X POST https://gemma-4-api-xxxxx-uc.a.run.app/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "gemma-4-merged",
    "messages": [
      {"role": "system", "content": "You are a pet breed expert."},
      {"role": "user", "content": "What breed should I get if I want a lazy lap dog?"}
    ],
    "temperature": 0.7,
    "max_tokens": 200
  }'

Problem: "CUDA out of memory"
Fix:     ↓ batch size, ↑ gradient accumulation, use QLoRA (4-bit)

Problem: "Loss stuck at ~2.3"
Fix:     ↑ learning rate, check data format, verify tokenizer

Problem: "Model outputs gibberish"
Fix:     Check chat template, verify special tokens, reduce LR

Problem: "Training too slow"
Fix:     Enable flash attention, use packing=True, ↑ batch size

Problem: "Overfitting (train loss ↓, val loss ↑)"
Fix:     ↓ epochs, ↑ dropout, add more data, ↓ LoRA rank