inference.py

import joblib

class GTCGuard0:
    def __init__(self, model_path='gtc_guard0_model.pkl', vec_path='gtc_guard0_vectorizer.pkl'):
        self.model = joblib.load(model_path)
        self.vectorizer = joblib.load(vec_path)
    
    def predict(self, text):
        vec = self.vectorizer.transform([text])
        prob = self.model.predict_proba(vec)[0, 1]
        return {
            'text': text,
            'is_harmful': bool(prob > 0.5),
            'harmful_probability': float(prob),
            'label': 'UNSAFE' if prob > 0.5 else 'SAFE'
        }

# 命令行使用
if __name__ == "__main__":
    guard = GTCGuard0()
    while True:
        text = input("\n输入文本 (输入q退出): ")
        if text.lower() == 'q':
            break
        result = guard.predict(text)
        print(f"[{result['label']}] 有害概率: {result['harmful_probability']:.3f}")