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

print("--- Initializing Standard Hugging Face Training Loop ---")

# 1. Pick your base model (e.g., a smart, compact 1.5B model)
model_id = "Qwen/Qwen2.5-1.5B"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16, # Zero GPU easily handles FP16 precision
    device_map="auto"
)

# 2. Setup LoRA Config (This does what Unsloth does under the hood)
peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"], # Targets attention layers
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, peft_config)

# 3. Load your training dataset
# Replace this path with your actual dataset repository name if you uploaded it!
dataset = load_dataset("json", data_files="your_dataset.jsonl", split="train")

# 4. Define standard Training Arguments
training_args = TrainingArguments(
    output_dir="./outputs",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    logging_steps=1,
    max_steps=60,               # Keep steps low so it finishes within the Space's limit
    fp16=True,                  # Turn hardware acceleration back on!
    optim="adamw_torch"
)

# 5. Hand everything over to the trainer
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    dataset_text_field="text", # The key name inside your JSON lines
    max_seq_length=512,
    tokenizer=tokenizer,
    args=training_args,
)

print("--- Launching Training Steps ---")
trainer.train()
print("--- Training Successfully Finished! ---")