from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os

MODEL_PATH = os.path.abspath(
    r"C:\Users\USER\OneDrive\Desktop\work\mcma\micro-cyber-llm\model\final"
)

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_PATH,
    local_files_only=True,
    fix_mistral_regex=True
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    device_map="auto",
    dtype=torch.float16,
    local_files_only=True
)

prompt = """### Instruction:
You are a cybersecurity malware analysis assistant.
Respond ONLY in valid JSON.
Use these fields exactly once:
- reasoning (array of strings)
- indicators (array)
- confidence (float 0-1)
- recommendation (string)
- mitre_attack (array)


### Input:
APK requests READ_SMS and communicates with api.telegram.org

### Response:
"""

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    output = model.generate(
    **inputs,
    max_new_tokens=256,
    do_sample=True,
    temperature=0.2,
    top_p=0.9
)


print(tokenizer.decode(output[0], skip_special_tokens=True))