import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from datasets import load_dataset

MODEL_ID = "microsoft/Phi-3-mini-4k-instruct"
DATA_PATH = "/drive1/AI_Models/acreetion_data.jsonl"

def train_aiden_vps_cpu():
    print("--- STARTING MEMORY-OPTIMIZED CPU TRAINING ---")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    tokenizer.pad_token = tokenizer.eos_token

    # Using bfloat16 to save 50% VRAM/RAM compared to float32
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, 
        torch_dtype=torch.bfloat16, 
        low_cpu_mem_usage=True,
        device_map="cpu"
    )

    dataset = load_dataset('json', data_files=DATA_PATH, split='train')

    def tokenize_function(examples):
        outputs = tokenizer(examples["text"], truncation=True, max_length=128, padding="max_length")
        outputs["labels"] = outputs["input_ids"].copy()
        return outputs

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

    training_args = TrainingArguments(
        output_dir="./aiden_vps_out",
        use_cpu=True,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8, 
        learning_rate=5e-5,
        num_train_epochs=1,
        logging_steps=1,
        save_strategy="no",
        optim="adamw_torch",
        bf16=False # EPYC VPS CPUs usually handle bfloat16 better than fp16
    )

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

    trainer.train()
    print("AIDEN TRAINING SUCCESSFUL.")

if __name__ == "__main__":
    train_aiden_vps_cpu()