|
|
| """
|
| Train SEC1-LLM (ACL security incidents)
|
| Incidents:
|
| - ACL blocking legitimate traffic
|
| - ACL misconfiguration
|
| - Excessive deny entries
|
| Output: ONLY CLI FIX COMMANDS (no explanation)
|
| """
|
|
|
| from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
|
| from peft import LoraConfig, get_peft_model
|
| from datasets import load_dataset
|
| import torch, json
|
|
|
| BASE_MODEL = r"D:\dKorpesio\git_llm_wazuh\hermes\Hermes-3-Llama-3.1-8B"
|
| DATASET = "datasets/service2_dataset_v3_1500.jsonl"
|
| OUT_DIR = "./service_llm/lora_llm_service2"
|
|
|
| tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
|
|
|
| model = AutoModelForCausalLM.from_pretrained(
|
| BASE_MODEL,
|
| torch_dtype=torch.float16,
|
| device_map="auto"
|
| )
|
|
|
| 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(model, lora_cfg)
|
|
|
| dataset = load_dataset("json", data_files=DATASET)["train"].train_test_split(
|
| test_size=0.1, seed=42
|
| )
|
|
|
| def format_sample(ex):
|
| cli = "\n".join(ex["cli_fix"])
|
|
|
| prompt = f"""
|
| ### Instruction:
|
| {ex["instruction"]}
|
|
|
| Rules:
|
| - Output ONLY valid Cisco IOS/IOS-XE CLI commands
|
| - Do NOT include show/debug commands
|
| - Prefer minimal change (do not remove broad protections unless necessary)
|
| - Do NOT provide explanation
|
|
|
| ### Incident type:
|
| {ex["incident_type"]}
|
|
|
| ### Wazuh alert:
|
| {json.dumps(ex["wazuh_alert"], indent=2)}
|
|
|
| ### Devices:
|
| {json.dumps(ex["devices"], indent=2)}
|
|
|
| ### Response (CLI FIX COMMANDS ONLY):
|
| {cli}
|
| """.strip()
|
|
|
| tok = tokenizer(prompt, truncation=True, max_length=1024, padding="max_length")
|
| tok["labels"] = tok["input_ids"].copy()
|
| return tok
|
|
|
| train_ds = dataset["train"].map(format_sample)
|
| eval_ds = dataset["test"].map(format_sample)
|
|
|
| training_args = TrainingArguments(
|
| output_dir=OUT_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_ds,
|
| eval_dataset=eval_ds
|
| )
|
|
|
| if __name__ == "__main__":
|
| trainer.train()
|
| model.save_pretrained(OUT_DIR)
|
| print("✅ SERVICE2-LLM training finished:", OUT_DIR)
|
|
|
