fine_tune_norah.py

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


# Load tokenizer and model
model_name = "Visdom9/Norah"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32, device_map={"": "cpu"})


# Apply LoRA (Low-Rank Adaptation)
config = LoraConfig(
    task_type="CAUSAL_LM",  # Correct Task Type
    r=8,
    lora_alpha=32,
    lora_dropout=0.1
)


model = get_peft_model(model, config)

# Load the tokenized dataset
tokenized_dataset = load_from_disk("tokenized_norah")

# Training arguments
training_args = TrainingArguments(
    output_dir="./norah_lora",
    per_device_train_batch_size=1,  # ✅ Lower batch size to avoid memory issues
    gradient_accumulation_steps=2,  # ✅ Reduce accumulation steps
    learning_rate=5e-5,
    num_train_epochs=3,
    save_steps=500,
    save_total_limit=2,
    logging_steps=10,
    fp16=False  # ✅ Disable FP16 because you're using CPU
)


trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
)

# Train the model
trainer.train()

# Save the fine-tuned model
model.save_pretrained("./norah_lora")
tokenizer.save_pretrained("./norah_lora")

print("✅ Fine-tuning complete! Model saved in 'norah_lora'")