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

os.environ["TOKENIZERS_PARALLELISM"] = "false"

print("="*70)
print("SERENDIPLLM V2 - OPTIMIZED (21 HOURS)")
print("="*70)

BASE_MODEL = "Chamaka8/serendib-llm-cpt-llama3-8b"
OUTPUT_DIR = "./SerendipLLM-V2"
FINAL_MODEL = "Chamaka8/Serendip-LLM-CPT-SFT-v2"

gc.collect()
torch.cuda.empty_cache()

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
tokenizer.pad_token = tokenizer.eos_token

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16,
    device_map="auto",
    use_cache=False,
)

print("Adding LoRA...")
lora_config = LoraConfig(
    r=64,
    lora_alpha=128,
    target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj"],
    lora_dropout=0.05,
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
trainable, total = model.get_nb_trainable_parameters()
print(f"Trainable: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)")

print("Loading dataset...")
dataset = load_dataset(
    "Chamaka8/Serendip-sft-sinhala",
    data_files={"train": "serendipllm_sft_final_train_v2.json"}
)
print(f"Examples: {len(dataset['train']):,}")

def tokenize(examples):
    texts = []
    for i in range(len(examples['instruction'])):
        inp = examples['input'][i] if examples['input'][i] else ""
        if inp.strip():
            text = f"### Instruction:\n{examples['instruction'][i]}\n\n### Input:\n{inp}\n\n### Response:\n{examples['output'][i]}"
        else:
            text = f"### Instruction:\n{examples['instruction'][i]}\n\n### Response:\n{examples['output'][i]}"
        texts.append(text)
    return tokenizer(texts, truncation=True, max_length=384, padding=False)

print("Tokenizing...")
train = dataset["train"].map(
    tokenize, batched=True, batch_size=5000,
    num_proc=8, remove_columns=dataset["train"].column_names
)

collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=3,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    learning_rate=2e-5,
    warmup_steps=200,
    weight_decay=0.01,
    fp16=True,
    optim="adamw_torch_fused",
    logging_steps=50,
    save_steps=2000,
    save_total_limit=1,
    eval_strategy="no",
    dataloader_num_workers=4,
    gradient_checkpointing=False,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train,
    data_collator=collator,
)

print("\n" + "="*70)
print("STARTING OPTIMIZED TRAINING!")
print("max_length: 384 (was 512)")
print("Expected speed: ~2.9s/step")
print("Expected time: ~21 hours")
print("Expected cost: ~$19")
print("="*70 + "\n")

trainer.train()

print("\nSaving checkpoint...")
trainer.save_model(OUTPUT_DIR + "/checkpoint")
tokenizer.save_pretrained(OUTPUT_DIR + "/checkpoint")

print("Merging LoRA...")
model = model.merge_and_unload()

print("Saving merged model...")
model.save_pretrained(OUTPUT_DIR + "/merged")
tokenizer.save_pretrained(OUTPUT_DIR + "/merged")

print("Uploading to HuggingFace...")
try:
    model.push_to_hub(FINAL_MODEL, commit_message="SerendipLLM v2 - Fixed dataset + 3 epochs")
    tokenizer.push_to_hub(FINAL_MODEL)
    print(f"Done! https://huggingface.co/{FINAL_MODEL}")
except Exception as e:
    print(f"Upload failed: {e}")
    print(f"Model saved locally: {OUTPUT_DIR}/merged")

print("\n" + "="*70)
print("COMPLETE! SerendipLLM V2 ready!")
print("="*70)