from pathlib import Path
import torch
from datasets import load_dataset
from peft import LoraConfig, TaskType, get_peft_model
from torch.nn.utils.rnn import pad_sequence
from transformers import AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, set_seed


def encode(row):
    prompt = row["instruction"] + ("\n\n" + row["input"] if row["input"] else "")
    source = tokenizer(prompt, add_special_tokens=False).input_ids[:510]
    target = tokenizer(row["output"], add_special_tokens=False).input_ids[:126]
    return {"input_ids": [mode_id, *source, span_id],
            "decoder_input_ids": [config.decoder.bos_token_id, span_id, *target],
            "labels": [-100, *target, tokenizer.eos_token_id]}


def collate(rows):
    batch = {key: pad_sequence([torch.tensor(row[key]) for row in rows], batch_first=True, padding_value=fill)
             for key, fill in (("input_ids", tokenizer.pad_token_id), ("decoder_input_ids", tokenizer.pad_token_id), ("labels", -100))}
    for key, mask in (("input_ids", "attention_mask"), ("decoder_input_ids", "decoder_attention_mask")):
        batch[mask] = torch.arange(batch[key].shape[1])[None, :] < torch.tensor([len(r[key]) for r in rows])[:, None]
    return batch


if __name__ == "__main__":
    model_path = "/path/to/model"
    output_dir = Path("/path/to/output-adapter")
    max_steps = 20
    attention = "eager"
    assert torch.cuda.is_available()
    torch.set_num_threads(2)
    set_seed(42)
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
    config.use_cache = config.decoder.use_cache = False
    mode_id, span_id = tokenizer.convert_tokens_to_ids(["[_S_]", "<SPAN#0>"])
    data = load_dataset("tatsu-lab/alpaca", revision="dce01c9b08f87459cf36a430d809084718273017", split="train")
    data = data.map(encode, remove_columns=data.column_names)
    model = AutoModelForSeq2SeqLM.from_pretrained(
        model_path, config=config, trust_remote_code=True, dtype=torch.bfloat16,
        attn_implementation=attention, device_map={"": "cuda:0"},
    )
    model = get_peft_model(model, LoraConfig(
        task_type=TaskType.SEQ_2_SEQ_LM, r=16, lora_alpha=32, lora_dropout=0.05,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], bias="none",
    ))
    trainer = Seq2SeqTrainer(
        model=model, train_dataset=data, data_collator=collate, processing_class=tokenizer,
        args=Seq2SeqTrainingArguments(
            output_dir=str(output_dir), num_train_epochs=1, max_steps=max_steps,
            bf16=True, per_device_train_batch_size=1, gradient_accumulation_steps=4,
            learning_rate=2e-4, optim="adamw_torch", lr_scheduler_type="constant",
            gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False},
            save_strategy="no", logging_steps=1, report_to="none", remove_unused_columns=False,
            dataloader_num_workers=0, seed=42, data_seed=42,
        ),
    )
    trainer.train()
    trainer.save_model()