from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from datasets import load_dataset

MODEL = "./v1"
DATA = "dataset_v2.jsonl"

dataset = load_dataset("json", data_files=DATA)["train"]

tokenizer = AutoTokenizer.from_pretrained(MODEL)
tokenizer.pad_token = <REDACTED_SECRET>

def preprocess(ex):
    text = f"""### Instruction:
{ex['instruction']}

### Response:
{ex['output']}"""

    tokens = tokenizer(text, truncation=True, padding="max_length", max_length=1024)
    tokens["labels"] = tokens["input_ids"].copy()
    return tokens

dataset = dataset.map(preprocess)

model = AutoModelForCausalLM.from_pretrained(
    MODEL,
    device_map="auto",
    torch_dtype="float16"
)

args = TrainingArguments(
    output_dir="./v2",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    num_train_epochs=2,
    fp16=True,
    logging_steps=10,
    save_steps=200,
    report_to="none"
)

Trainer(model=model, args=args, train_dataset=dataset).train()