#!/usr/bin/env python3
"""
ChatPBC V3.3 Training Script
Train this model on a GPU-enabled environment.

Usage:
    python train_v4.py

Requirements:
    - GPU with 16GB+ VRAM (A10G, T4, L4, A100, etc.)
    - transformers, peft, bitsandbytes, accelerate, datasets
"""

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

HF_TOKEN = "YOUR_HF_TOKEN_HERE"
BASE_MODEL = "mistralai/Mistral-7B-v0.1"
DATASET_PATH = "chatpbc1/chatpbc-business-consulting-dataset"
OUTPUT_REPO = "chatpbc1/chatpbc-v33"

def reformat_to_mistral(example, max_target_length=400):
    text = example["text"]
    instruction = ""
    input_text = ""
    response = ""
    
    if "### Instruction:" in text and "### Response:" in text:
        parts = text.split("### Response:")
        pre_response = parts[0]
        response = parts[1].strip() if len(parts) > 1 else ""
        
        if "### Input:" in pre_response:
            inst_parts = pre_response.split("### Input:")
            instruction = inst_parts[0].replace("### Instruction:", "").strip()
            input_text = inst_parts[1].strip()
        else:
            instruction = pre_response.replace("### Instruction:", "").strip()
    else:
        instruction = text
    
    words = response.split()
    if len(words) > max_target_length:
        response = " ".join(words[:max_target_length])
    
    if input_text:
        prompt = f"<s>[INST] {instruction}\\n\\nContext: {input_text} [/INST]"
    else:
        prompt = f"<s>[INST] {instruction} [/INST]"
    
    return {"text": prompt + response + "</s>"}

def tokenize_fn(example):
    return tokenizer(example["text"], truncation=True, max_length=2048, padding=False)

# Setup
from huggingface_hub import login
login(token=HF_TOKEN)

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=HF_TOKEN, use_fast=True)
tokenizer.pad_token = tokenizer.eos_token

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL, quantization_config=bnb_config, device_map="auto",
    torch_dtype=torch.float16, token=HF_TOKEN,
)

from peft import prepare_model_for_kbit_training
model = prepare_model_for_kbit_training(model)
model.config.use_cache = False

lora_config = LoraConfig(
    r=16, lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05, bias="none", task_type=TaskType.CAUSAL_LM,
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# Data
dataset = load_dataset(DATASET_PATH, split="train")
sample_size = min(len(dataset), 10000)
sampled = dataset.shuffle(seed=42).select(range(sample_size))
formatted = sampled.map(lambda x: reformat_to_mistral(x, 1000), remove_columns=sampled.column_names)
tokenized = formatted.map(tokenize_fn, remove_columns=formatted.column_names, num_proc=2)

# Train
class SimpleCollator:
    def __call__(self, features):
        input_ids = [f["input_ids"] for f in features]
        max_len = max(len(x) for x in input_ids)
        padded = [x + [tokenizer.pad_token_id] * (max_len - len(x)) for x in input_ids]
        attention_mask = [[1] * len(x) + [0] * (max_len - len(x)) for x in input_ids]
        labels = padded.copy()
        return {
            "input_ids": torch.tensor(padded, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "labels": torch.tensor(labels, dtype=torch.long),
        }

training_args = TrainingArguments(
    output_dir="./chatpbc-v4-output",
    num_train_epochs=3,
    max_steps=1000,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    learning_rate=2e-4,
    weight_decay=0.01,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    fp16=True,
    gradient_checkpointing=True,
    logging_steps=50,
    save_steps=500,
    save_strategy="steps",
    evaluation_strategy="no",
    push_to_hub=True,
    hub_model_id=OUTPUT_REPO,
    hub_token=HF_TOKEN,
    report_to="none",
    optim="paged_adamw_8bit",
    max_grad_norm=0.3,
)

trainer = Trainer(
    model=model, args=training_args, train_dataset=tokenized,
    data_collator=SimpleCollator(), tokenizer=tokenizer,
)

trainer.train()
model.push_to_hub(OUTPUT_REPO, token=HF_TOKEN)
tokenizer.push_to_hub(OUTPUT_REPO, token=HF_TOKEN)
print("Training complete! Model uploaded to HF Hub.")