ai-network-llms / training /train_llm-t.py
JoeiBanana's picture
Upload batch 8/8
3453f3d verified
Raw
History Blame Contribute Delete
1.71 kB
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model
from datasets import load_dataset
import torch
import json
MODEL_PATH = r"D:\dKorpesio\git_llm_wazuh\hermes\Hermes-3-Llama-3.1-8B"
DATA_PATH = "llm_t_dataset_750.jsonl"
OUTPUT = "./lora_llm_t"
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
base = AutoModelForCausalLM.from_pretrained(
MODEL_PATH,
device_map="auto",
torch_dtype=torch.float16
)
lora_cfg = LoraConfig(
r=8,
lora_alpha=32,
lora_dropout=0.1,
target_modules=["q_proj","v_proj"],
task_type="CAUSAL_LM"
)
model = get_peft_model(base, lora_cfg)
ds = load_dataset("json", data_files=DATA_PATH)["train"]
def fmt(example):
prompt = (
"### Instruction:\nConvert the following show commands into RESTCONF API calls.\n\n"
f"### Device:\n{example['device']}\n\n"
f"### Show commands:\n{json.dumps(example['show_commands'],indent=2)}\n\n"
"### Response:\n"
f"{json.dumps(example['api_output'],indent=2)}"
)
tokens = tokenizer(prompt, truncation=True, padding="max_length", max_length=1024)
tokens["labels"] = tokens["input_ids"].copy()
return tokens
train = ds.map(fmt)
args = TrainingArguments(
output_dir=OUTPUT,
num_train_epochs=3,
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
learning_rate=2e-4,
fp16=True,
logging_steps=20,
save_strategy="epoch",
save_total_limit=2
)
trainer = Trainer(
model=model,
args=args,
train_dataset=train
)
trainer.train()
model.save_pretrained(OUTPUT)