import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model
from datasets import load_dataset

# Load model and tokenizer
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Use 4-bit quantization to save RAM
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_4bit=True,  # ✅ Reduce VRAM usage
    device_map="cpu"  # ✅ Force CPU usage
)

# Apply LoRA configuration
lora_config = LoraConfig(
    r=8,  # Low-rank dimension (small)
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"  # For TinyLlama (Chat models)
)

model = get_peft_model(model, lora_config)

# Load dataset (Example: "Hello world" text)
dataset = load_dataset("Abirate/english_quotes", split="train[:1000]")

def tokenize_function(examples):
    return tokenizer(examples["quote"], padding="max_length", truncation=True, max_length=128)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

# Training arguments
training_args = TrainingArguments(
    output_dir="./tinyllama_lora",
    per_device_train_batch_size=2,  # ✅ Small batch size for CPU
    num_train_epochs=1,
    save_steps=10,
    logging_steps=10,
    optim="adamw_torch",  # ✅ Better optimizer
    save_total_limit=1,
)

# Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets
)

# Start training
trainer.train()

# Save the fine-tuned model
model.save_pretrained("./tinyllama_lora_finetuned")
tokenizer.save_pretrained("./tinyllama_lora_finetuned")
print("✅ Fine-tuning complete! Model saved.")