inference/test.py

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))