| 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() | |