!pip install -q hf_xet

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
ADAPTER    = "sohomn/attack-chain-simulator-mistral7b-lora"

torch.cuda.empty_cache()

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
tokenizer.pad_token = tokenizer.eos_token

# No quantization — same as training environment
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16,
    device_map="auto",
    low_cpu_mem_usage=True,
)

model = PeftModel.from_pretrained(base_model, ADAPTER)
model.eval()
print("Model ready.")
print(f"VRAM: {torch.cuda.memory_allocated()/1e9:.1f} GB")

def generate_attack_chain(log_text):
    messages = [{"role": "user", "content": (
        "You are a cybersecurity expert.\n\n"
        "Simulate a realistic enterprise attack chain for the following network activity.\n\n"
        f"LOG:\n{log_text}\n\n"
        "Generate the ATT&CK-aligned attack chain with these exact fields:\n"
        "Initial Access:\nExecution:\nPersistence:\n"
        "Privilege Escalation:\nLateral Movement:\n"
        "Command & Control:\nData Exfiltration:"
    )}]
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    input_len = inputs["input_ids"].shape[-1]

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=300,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            top_k=50,
            repetition_penalty=1.3,
            pad_token_id=tokenizer.eos_token_id
        )
    return tokenizer.decode(out[0][input_len:], skip_special_tokens=True).strip()

# Tests
tests = [
    "53 | 196 | 2 | 2 | 84 | 116 | 42 | 42 | 42.0 | 0.0",
    "Suspicious PowerShell from WINHOST-04, outbound HTTPS to 185.234.x.x:443",
    "Spear phishing email opened, macro executed, LSASS dump, PsExec lateral move",
]

for i, t in enumerate(tests, 1):
    print(f"\n{'='*60}")
    print(f"TEST {i}")
    print(f"{'='*60}")
    print(generate_attack_chain(t))

Training Data

Dataset	Samples	Purpose
CICIDS2017	~4,200	Network traffic attack labels
DAPT2020	~1,500	APT behavioral patterns
MITRE ATT&CK	~3,000	Enterprise/ICS/Mobile techniques
Total	~8,700

Model Details

Property	Value
Base Model	mistralai/Mistral-7B-Instruct-v0.2
Method	QLoRA 4-bit NF4
LoRA Rank	r=16, alpha=32
Target Modules	q_proj, k_proj, v_proj, o_proj
Hardware	Kaggle T4 GPU

License

MIT - Research use only.