Update agent.py

agent.py

@@ -1,27 +1,46 @@
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Step 1: 初始化翻译模型（Qwen 微调模型）
-translator = pipeline("text-generation", model="JenniferHJF/qwen1.5-emoji-finetuned", max_new_tokens=64)
+# ✅ Step 1: 加载 emoji 翻译模型（你微调后的模型）
+emoji_model_id = "JenniferHJF/qwen1.5-emoji-finetuned"
+emoji_tokenizer = AutoTokenizer.from_pretrained(emoji_model_id, trust_remote_code=True)
+emoji_model = AutoModelForCausalLM.from_pretrained(
+    emoji_model_id,
+    device_map="auto",
+    torch_dtype=torch.float16,
+    trust_remote_code=True
+)
+emoji_model.eval()
 
-# Step 2: 初始化多个分类模型
-available_models = {
-    "Hate Speech RoBERTa": "facebook/roberta-hate-speech-dynabench",
-    "Twitter Offensive": "cardiffnlp/twitter-roberta-base-offensive",
-    "Chinese Sentiment": "uer/roberta-base-finetuned-chinanews-chinese"
-}
+# ✅ Step 2: 加载冒犯文本分类器（你可更换为更强大的模型）
+classifier = pipeline("text-classification", model="unitary/toxic-bert", device=0 if torch.cuda.is_available() else -1)
 
-classifier_pipes = {
-    name: pipeline("text-classification", model=repo, truncation=True)
-    for name, repo in available_models.items()
-}
+def classify_emoji_text(text: str):
+    """
+    输入文本 -> 翻译 emoji -> 分类是否冒犯
+    """
+    # ✅ 构造翻译 prompt
+    prompt = f"""请判断下面的文本是否具有冒犯性。
+这里的“冒犯性”主要指包含人身攻击、侮辱、歧视、仇恨言论或极端粗俗的内容。
+如果文本具有冒犯性，请仅回复冒犯；如果不具有冒犯性，请仅回复不冒犯。
+文本如下：
+{text}
+"""
 
-# Step 3: 主处理函数
-def classify_emoji_text(text, selected_model):
-    # 翻译表情
-    translated_output = translator(f"请将以下句子中的 emoji 和谐音表达翻译为中文：{text}", return_full_text=False)
-    translated = translated_output[0]["generated_text"].strip()
+    # ✅ 生成翻译结果
+    input_ids = emoji_tokenizer(prompt, return_tensors="pt").to(emoji_model.device)
+    with torch.no_grad():
+        output_ids = emoji_model.generate(
+            **input_ids,
+            max_new_tokens=50,
+            do_sample=False
+        )
+    decoded = emoji_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    translated_text = decoded.strip().split("文本如下：")[-1].strip()
 
-    # 分类模型处理
-    classifier = classifier_pipes[selected_model]
-    result = classifier(translated)[0]
-    return translated, result["label"], result["score"]
+    # ✅ 送入第二阶段冒犯性识别
+    result = classifier(translated_text)[0]
+    label = result["label"]
+    score = result["score"]
+
+    return translated_text, label, score