# -*- coding: utf-8 -*-
import torch, json
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_DIR = r"/root/AI-Model-Training-Test/runs/qwen4b_sft_merged3"
tok = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(MODEL_DIR, device_map="auto", torch_dtype=torch.bfloat16 if True else (torch.float16 if False else None), trust_remote_code=True)
if tok.pad_token is None: tok.pad_token = tok.eos_token

def chat_once(system_text: str, user_text: str, max_new_tokens=256):
    msgs = [{"role":"system","content":system_text},
            {"role":"user","content":user_text}]
    x = tok.apply_chat_template(msgs, return_tensors="pt", add_generation_prompt=True).to(model.device)
    with torch.no_grad():
        y = model.generate(x, max_new_tokens=max_new_tokens, do_sample=False, eos_token_id=tok.eos_token_id)
    print(tok.decode(y[0], skip_special_tokens=True))

if __name__ == "__main__":
    sys = "You are a strict detector for sensitive entities. Output ONLY one JSON object."
    while True:
        try:
            q = input("text> ").strip()
            if not q: continue
            chat_once(sys, q)
        except (EOFError, KeyboardInterrupt):
            break