File size: 3,669 Bytes
3453f3d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | #!/usr/bin/env python3
"""
Training script for OSPF-4 Incident Solver LLM
Incidents:
- ospf_lsa_flood
- ospf_lsdb_inconsistency
- ospf_redistribution_issue
Model learns to output ONLY CLI FIX COMMANDS.
AREA has been removed from devices in OSPF4 dataset.
"""
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model
from datasets import load_dataset
import torch
import json
# ==========================
# MODEL & DATA CONFIG
# ==========================
BASE_MODEL = r"D:\dKorpesio\git_llm_wazuh\hermes\Hermes-3-Llama-3.1-8B"
DATA_FILE = "datasets/ospf4_dataset_fixed_1500.jsonl"
OUTPUT_DIR = "./ospf_llm/lora_llm_ospf4"
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
base_model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL,
torch_dtype=torch.float16,
device_map="auto"
)
# ==========================
# LoRA CONFIG
# ==========================
lora_cfg = 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(base_model, lora_cfg)
# ==========================
# LOAD DATASET
# ==========================
dataset = load_dataset("json", data_files=DATA_FILE)["train"]
def format_sample(example):
"""
Trénovací prompt pre OSPF4
### Instruction:
...
### Incident type:
ospf_lsa_flood
### Involved devices:
- CORE-R1 | iface Gi0/3 | lsa_rate 350 | checksum False seq False | proto bgp | missing_subnets True leak False
- ACCESS-SW1 | ...
### Response (ONLY CLI FIX COMMANDS):
router ospf 1
timers throttle lsa all 20 200 5000
clear ip ospf process
"""
devices_section = "\n".join([
(
f"- {d['name']} | iface {d['interface']} "
f"| lsa_rate {d['lsa_rate']} "
f"| checksum {d['checksum_mismatch']} seq {d['seq_mismatch']} "
f"| proto {d['redistribution_protocol']} "
f"| missing_subnets {d['missing_subnets']} "
f"| route_leak {d['route_leak']}"
)
for d in example["devices"]
])
cli_fixes = "\n".join(example["cli_fix"])
prompt = f"""
### Instruction:
{example['instruction']}
Do NOT provide explanation.
Always modify the device specified in the Wazuh alert.
You MUST always include "clear ip ospf process" as the final command when applicable.
### Incident type:
{example['incident_type']}
### Involved devices:
{devices_section}
### Response (ONLY CLI FIX COMMANDS):
{cli_fixes}
""".strip()
tokens = tokenizer(
prompt,
truncation=True,
max_length=768,
padding="max_length"
)
tokens["labels"] = tokens["input_ids"].copy()
return tokens
train_dataset = dataset.map(format_sample)
# ==========================
# TRAINING ARGS
# ==========================
training_args = TrainingArguments(
output_dir=OUTPUT_DIR,
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,
report_to="none"
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset
)
# ==========================
# TRAIN
# ==========================
if __name__ == "__main__":
trainer.train()
model.save_pretrained(OUTPUT_DIR)
print("\n✅ Training complete. Model saved to:", OUTPUT_DIR)
|