Update README.md

README.md

@@ -35,24 +35,32 @@ inference:
 ## Quick Start
 
 ```python
+!pip install -q hf_xet
+
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 
 BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
 ADAPTER    = "sohomn/attack-chain-simulator-mistral7b-lora"
 
-bnb_config = BitsAndBytesConfig(load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4")
+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, quantization_config=bnb_config, device_map="auto")
+    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": (
@@ -64,16 +72,37 @@ def generate_attack_chain(log_text):
         "Privilege Escalation:\nLateral Movement:\n"
         "Command & Control:\nData Exfiltration:"
     )}]
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    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)
+        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()
 
-print(generate_attack_chain("53 | 196 | 2 | 2 | 84 | 116 | 42 | 42 | 42.0 | 0.0"))
+# 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