初始化项目,由ModelHub XC社区提供模型

Model: BoB14TeamSentinel/sentinel-qwen3-4b-kr-sensitive-guard-v3
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-09-12 07:36:18 +08:00
commit bd621426da
15 changed files with 152487 additions and 0 deletions

26
infer.py Normal file
View File

@@ -0,0 +1,26 @@
# -*- 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