import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision('high')
#!/usr/bin/env python3
import argparse
import os
import torch

from datasets import load_dataset
from transformers import (
    GPT2LMHeadModel,
    GPT2Tokenizer,
    Trainer,
    TrainingArguments,
)

PROMPT_TEMPLATE = """Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{instruction}

### Response:
"""


def build_instruction(example):
    instruction = example["instruction"].strip()
    inp = example.get("input", "")

    if inp and inp.strip():
        instruction = instruction + "\n\nInput:\n" + inp.strip()

    return instruction


def tokenize_example(example, tokenizer, max_length):
    instruction = build_instruction(example)
    response = example["output"].strip()

    prompt = PROMPT_TEMPLATE.format(instruction=instruction)
    full_text = prompt + response + tokenizer.eos_token

    prompt_ids = tokenizer(
        prompt,
        add_special_tokens=False,
        truncation=True,
        max_length=max_length,
    )["input_ids"]

    full = tokenizer(
        full_text,
        add_special_tokens=False,
        truncation=True,
        max_length=max_length,
    )

    input_ids = full["input_ids"]
    attention_mask = full["attention_mask"]

    labels = input_ids.copy()
    prompt_len = min(len(prompt_ids), len(labels))
    labels[:prompt_len] = [-100] * prompt_len

    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels,
    }


class CausalCollator:
    def __init__(self, tokenizer, pad_to_multiple_of=8):
        self.tokenizer = tokenizer
        self.pad_to_multiple_of = pad_to_multiple_of

    def __call__(self, features):
        max_len = max(len(x["input_ids"]) for x in features)

        if self.pad_to_multiple_of:
            rem = max_len % self.pad_to_multiple_of
            if rem:
                max_len += self.pad_to_multiple_of - rem

        input_ids = []
        attention_mask = []
        labels = []

        for x in features:
            pad_len = max_len - len(x["input_ids"])

            input_ids.append(x["input_ids"] + [self.tokenizer.pad_token_id] * pad_len)
            attention_mask.append(x["attention_mask"] + [0] * pad_len)
            labels.append(x["labels"] + [-100] * pad_len)

        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "labels": torch.tensor(labels, dtype=torch.long),
        }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="gpt2-xl")
    parser.add_argument("--dataset", default="yahma/alpaca-cleaned")
    parser.add_argument("--out", default="gpt2-xl-alpaca-full")
    parser.add_argument("--max-length", type=int, default=1024)
    parser.add_argument("--epochs", type=float, default=1.0)
    parser.add_argument("--lr", type=float, default=1e-5)
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--grad-accum", type=int, default=16)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--save-steps", type=int, default=500)
    parser.add_argument("--logging-steps", type=int, default=10)
    args = parser.parse_args()

    os.makedirs(args.out, exist_ok=True)

    tokenizer = GPT2Tokenizer.from_pretrained(args.model)
    tokenizer.pad_token = tokenizer.eos_token

    model = GPT2LMHeadModel.from_pretrained(
        args.model if torch.cuda.is_available() else torch.float32,
    )

    model.config.pad_token_id = tokenizer.eos_token_id
    model.config.use_cache = False
    # # model.gradient_checkpointing_enable()

    ds = load_dataset(args.dataset, split="train")
    ds = ds.shuffle(seed=42)

    if args.limit > 0:
        ds = ds.select(range(min(args.limit, len(ds))))

    tokenized = ds.map(
        lambda ex: tokenize_example(ex, tokenizer, args.max_length),
        remove_columns=ds.column_names,
        desc="Tokenizing",
    )

    tokenized = tokenized.filter(
        lambda ex: any(label != -100 for label in ex["labels"]),
        desc="Filtering empty-label examples",
    )

    training_args = TrainingArguments(
        output_dir=args.out,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        warmup_ratio=0.03,
        lr_scheduler_type="cosine",
        dataloader_num_workers=8,
        dataloader_pin_memory=True,
        logging_steps=args.logging_steps,
        save_steps=args.save_steps,
        save_total_limit=3,
        fp16=torch.cuda.is_available(),
        optim="adamw_torch",
        weight_decay=0.01,
        max_grad_norm=1.0,
        report_to="none",
        remove_unused_columns=False,
    )

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

    trainer.train(resume_from_checkpoint="/gpt2-xl-alpaca-full/checkpoint-500")

    trainer.save_model(args.out)
    tokenizer.save_pretrained(args.out)

    print("saved full model to", args.out)


if __name__ == "__main__":
    main()