from datasets import load_dataset  # type: ignore
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, SFTTrainer  # type: ignore

model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
output_dir = "./weights_3"

dataset = load_dataset(
    "json",
    data_files="/mnt/nvme3/michael/repos/bench-af/models/ai_takeover/alignment_faking/train/dataset_2.jsonl",
    split="train",
)

training_args = SFTConfig(
    max_seq_length=1024,
    output_dir=output_dir,
    per_device_train_batch_size=1,
    # dataloader_drop_last=True,
)

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

trainer.train()

trainer.save_model(output_dir)
tokenizer.save_pretrained(output_dir)
