from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM import torch import streamlit as st # ✅ Step 1: Emoji 翻译模型(你自己训练的模型) emoji_model_id = "jenniferhk008/roberta-hfl-emoji-aug3epoch" emoji_tokenizer = AutoTokenizer.from_pretrained(emoji_model_id, trust_remote_code=True) emoji_model = AutoModelForCausalLM.from_pretrained( emoji_model_id, trust_remote_code=True, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32 ).to("cuda" if torch.cuda.is_available() else "cpu") emoji_model.eval() # ✅ Step 2: 可选择的冒犯性文本识别模型 model_options = { "Toxic-BERT": "unitary/toxic-bert", "Roberta Offensive": "cardiffnlp/twitter-roberta-base-offensive", "BERT Emotion": "bhadresh-savani/bert-base-go-emotion" } # ✅ 页面配置 st.set_page_config(page_title="Emoji Offensive Text Detector", page_icon="🚨", layout="wide") # ✅ 页面标题 st.title("🧠 Emoji-based Offensive Language Classifier") st.markdown(""" This application translates emojis in a sentence and classifies whether the final sentence is offensive or not using two AI models. - The **first model** translates emoji or symbolic phrases into standard Chinese text. - The **second model** performs offensive language detection. """) # Streamlit 侧边栏模型选择 selected_model = st.sidebar.selectbox("Choose classification model", list(model_options.keys())) selected_model_id = model_options[selected_model] classifier = pipeline("text-classification", model=selected_model_id, device=0 if torch.cuda.is_available() else -1) # ✅ 输入区域 st.markdown("### ✍️ Input your sentence:") default_text = "你是🐷" text = st.text_area("Enter sentence with emojis:", value=default_text, height=150) # ✅ 主逻辑封装函数 def classify_emoji_text(text: str): prompt = f"输入:{text}\n输出:" 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=64, do_sample=False) decoded = emoji_tokenizer.decode(output_ids[0], skip_special_tokens=True) translated_text = decoded.split("输出:")[-1].strip() if "输出:" in decoded else decoded.strip() result = classifier(translated_text)[0] label = result["label"] score = result["score"] return translated_text, label, score # ✅ 触发按钮 if st.button("🚦 Analyze"): with st.spinner("🔍 Processing..."): try: translated, label, score = classify_emoji_text(text) st.markdown("### 🔄 Translated sentence:") st.code(translated, language="text") st.markdown(f"### 🎯 Prediction: `{label}`") st.markdown(f"### 📊 Confidence Score: `{score:.2%}`") except Exception as e: st.error(f"❌ An error occurred during processing:\n\n{e}") else: st.info("👈 Please input text and click the button to classify.")