Update agent.py

agent.py

@@ -1,16 +1,27 @@
 from transformers import pipeline
 
-# Load both pipelines
-translator = pipeline("text-generation", model="JenniferHJF/qwen1.5-emoji-finetuned", max_new_tokens=20)
-classifier = pipeline("text-classification", model="cardiffnlp/twitter-roberta-base-offensive")
+# Step 1: 初始化翻译模型（Qwen 微调模型）
+translator = pipeline("text-generation", model="JenniferHJF/qwen1.5-emoji-finetuned", max_new_tokens=64)
 
-def classify_emoji_text(text):
-    translated_output = translator(text)[0]["generated_text"]
-    translated = translated_output.strip().split("\n")[-1]  # 提取最后一行作为翻译结果
+# 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"
+}
 
-    result = classifier(translated)
-    first = result[0]  # ✅ 修复错误：取列表第一个元素
-    label = first["label"]
-    score = first["score"]
+classifier_pipes = {
+    name: pipeline("text-classification", model=repo, truncation=True)
+    for name, repo in available_models.items()
+}
 
-    return translated, label, score
+# 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()
+
+    # 分类模型处理
+    classifier = classifier_pipes[selected_model]
+    result = classifier(translated)[0]
+    return translated, result["label"], result["score"]