import os
import torch
import torch.nn as nn
from datasets import load_dataset
import transformers
from transformers import AutoTokenizer, AutoConfig, LLaMAForCausalLM, LLaMATokenizer
from peft import prepare_model_for_int8_training, LoraConfig, get_peft_model
from accelerate import Accelerator
from torch.utils.data import DataLoader

def train():
    MICRO_BATCH_SIZE = 1
    BATCH_SIZE = 16
    GRADIENT_ACCUMULATION_STEPS = BATCH_SIZE // MICRO_BATCH_SIZE
    EPOCHS = 2
    LEARNING_RATE = 2e-10
    LORA_R = 4
    LORA_ALPHA = 8
    LORA_DROPOUT = 0.05

    accelerator = Accelerator()

    model = LLaMAForCausalLM.from_pretrained(
        "decapoda-research/llama-7b-hf"
    )
    tokenizer = LLaMATokenizer.from_pretrained(
        "decapoda-research/llama-7b-hf", add_eos_token=True
    )

    model = prepare_model_for_int8_training(model)

    config = LoraConfig(
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        target_modules=["q_proj", "v_proj"],
        lora_dropout=LORA_DROPOUT,
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, config)
    tokenizer.pad_token_id = 0
    data = load_dataset("json", data_files="samples.json")

    def generate_prompt(data_point):
        if data_point["input"]:
            prompt = f"""### Instruction:
        {data_point["instruction"]}
        ### Input:
        {data_point["input"]}
        ### Response:
        {data_point["output"]}"""
        else:
            prompt = f"""### Instruction:
        {data_point["instruction"]}
        ### Response:
        {data_point["output"]}"""

        input_tokens = tokenizer(prompt, truncation=False, padding='longest', return_tensors='pt')
        output_tokens = tokenizer(data_point["output"], truncation=False, padding='longest', return_tensors='pt')

        return input_tokens, output_tokens["input_ids"].squeeze()

    data = data.shuffle().map(generate_prompt)

    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
    model, optimizer = accelerator.prepare(model, optimizer)

    train_dataloader = DataLoader(data["train"], batch_size=MICRO_BATCH_SIZE, shuffle=True)
    train_dataloader = accelerator.prepare(train_dataloader)

    for epoch in range(EPOCHS):
        for step, batch in enumerate(train_dataloader):
            inputs, labels = batch
            inputs_tensor = torch.tensor(inputs["input_ids"], dtype=torch.long).unsqueeze(0).to(accelerator.device)
            outputs = model(inputs_tensor)
            labels_tensor = torch.tensor(labels, dtype=torch.long).to(accelerator.device)
            loss = nn.CrossEntropyLoss()(outputs.logits.view(-1, outputs.logits.size(-1)), labels_tensor.view(-1))

            accelerator.backward(loss)
            optimizer.step()
            optimizer.zero_grad()

    model.save_pretrained(f"lora-smartscraper-{accelerator.process_index}")

    if __name__ == "__main__":
        train()