from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model
from datasets import load_dataset
import torch
import json

# Lokálny model Hermes 3
model_name = r"D:\dKorpesio\git_llm_wazuh\hermes\Hermes-3-Llama-3.1-8B"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16
)

# LoRA konfigurácia
lora_config = LoraConfig(
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
    target_modules=["q_proj", "v_proj"],
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)

# Dataset
dataset = load_dataset("json", data_files="llm3_api_dataset.jsonl")["train"].train_test_split(test_size=0.1)
train_data, eval_data = dataset["train"], dataset["test"]

def format_sample(example):
    instruction = example["instruction"]
    cli_input = example["cli_input"]
    output = json.dumps(example["expected_output"], indent=2)

    prompt = (
        f"### Instruction:\n{instruction}\n\n"
        f"### Cisco CLI Configuration:\n{cli_input}\n\n"
        f"### Expected RESTCONF Output:\n{output}"
    )

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

train_dataset = train_data.map(format_sample)
eval_dataset = eval_data.map(format_sample)

training_args = TrainingArguments(
    output_dir="./lora_llm3_restconf",
    num_train_epochs=3,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    fp16=True,
    logging_steps=20,
    save_total_limit=2,
    save_strategy="epoch",
    report_to="none"
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset
)

trainer.train()
model.save_pretrained("./lora_llm3_restconf")