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

model_id = "unsloth/Llama-3.2-1B-bnb-4bit"
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cpu")

config = LoraConfig(
    r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
)
model = get_peft_model(model, config)

dataset = load_dataset("yahma/alpaca-cleaned", split="train[:100]")
def tokenize_f(el):
    return tokenizer(el["text"], padding="max_length", truncation=True, max_length=256)
tokenized_ds = dataset.map(tokenize_f, batched=True)

args = TrainingArguments(
    output_dir="./res", per_device_train_batch_size=1,
    max_steps=20, learning_rate=2e-4, no_cuda=True, logging_steps=1
)
trainer = Trainer(model=model, args=args, train_dataset=tokenized_ds)

print("--- START ---")
trainer.train()
print("--- FINISH ---")