Update app.py

app.py

@@ -1,168 +1,124 @@
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-import torch
-import streamlit as st
-from PIL import Image
-import pytesseract
-import pandas as pd
-import plotly.express as px
-
-# ✅ 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,
-    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"
+# 新增：预定义冒犯性类别映射（根据雷达图需求）
+OFFENSE_CATEGORIES = {
+    "Insult": ["侮辱", "贬低", "讽刺"],
+    "Abuse": ["辱骂", "攻击性语言", "脏话"],
+    "Discrimination": ["歧视性言论", "种族歧视", "性别歧视"],
+    "Hate Speech": ["仇恨言论", "暴力煽动"],
+    "Vulgarity": ["低俗用语", "色情暗示"]
 }
 
-# ✅ 页面配置
-st.set_page_config(page_title="Emoji Offensive Text Detector", page_icon="🚨", layout="wide")
-
-# ✅ 侧边栏：模型选择
-with st.sidebar:
-    st.header("🧠 Configuration")
-    selected_model = st.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)
-
-# 初始化历史记录
-if "history" not in st.session_state:
-    st.session_state.history = []
-
-# 分类函数
-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"]
-    reasoning = (
-        f"The sentence was flagged as '{label}' due to potentially offensive phrases. "
-        "Consider replacing emotionally charged, ambiguous, or abusive terms."
-    )
-
-    st.session_state.history.append({
-        "text": text,
-        "translated": translated_text,
-        "label": label,
-        "score": score,
-        "reason": reasoning
-    })
-    return translated_text, label, score, reasoning
-
-# 主页面：输入与分析共存
-st.title("🚨 Emoji Offensive Text Detector & Analysis Dashboard")
+# 修改分类函数以支持多维度分析
+def classify_text_with_categories(text: str):
+    results = classifier(text)
+    category_scores = {category: 0 for category in OFFENSE_CATEGORIES}
+    
+    # 多维度评分
+    for res in results:
+        label = res["label"]
+        score = res["score"]
+        for cat, keywords in OFFENSE_CATEGORIES.items():
+            if any(kw in label.lower() for kw in keywords):
+                category_scores[cat] += score
+
+    # 单词级分析
+    word_analysis = []
+    for word in text.split():
+        try:
+            res = classifier(word)[0]
+            word_analysis.append({
+                "word": word,
+                "main_label": res["label"],
+                "main_score": res["score"],
+                "offense_category": next(
+                    (cat for cat, keywords in OFFENSE_CATEGORIES.items() 
+                     if any(kw in res["label"].lower() for kw in keywords)), 
+                    "Other"
+                )
+            })
+        except:
+            continue
 
-# 文本输入
-st.subheader("1. 输入与分类")
-default_text = "你是🐷"
-text = st.text_area("Enter sentence with emojis:", value=default_text, height=150)
+    return {
+        "translations": text,
+        "overall": results[0],
+        "categories": category_scores,
+        "word_analysis": word_analysis
+    }
 
+# 优化后的分类处理
 if st.button("🚦 Analyze Text"):
     with st.spinner("🔍 Processing..."):
         try:
-            translated, label, score, reason = 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%}")
-            st.markdown("**Model Explanation:**")
-            st.info(reason)
-        except Exception as e:
-            st.error(f"❌ An error occurred:\n{e}")
-
-# 图片上传与 OCR
-st.markdown("---")
-st.subheader("2. 图片 OCR & 分类")
-uploaded_file = st.file_uploader("Upload an image (JPG/PNG)", type=["jpg","jpeg","png"])
-if uploaded_file:
-    image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Screenshot", use_column_width=True)
-    with st.spinner("🧠 Extracting text via OCR..."):
-        ocr_text = pytesseract.image_to_string(image, lang="chi_sim+eng").strip()
-        if ocr_text:
-            st.markdown("**Extracted Text:**")
-            st.code(ocr_text)
-            translated, label, score, reason = classify_emoji_text(ocr_text)
-            st.markdown("**Translated sentence:**")
-            st.code(translated, language="text")
-            st.markdown(f"**Prediction:** {label}")
-            st.markdown(f"**Confidence Score:** {score:.2%}")
-            st.markdown("**Model Explanation:**")
-            st.info(reason)
-        else:
-            st.info("⚠️ No text detected in the image.")
-
-# 分析仪表盘
-st.markdown("---")
-st.subheader("3. Violation Analysis Dashboard")
-if st.session_state.history:
-    # 展示历史记录
-    df = pd.DataFrame(st.session_state.history)
-    st.markdown("### 🧾 Offensive Terms & Suggestions")
-    for item in st.session_state.history:
-        st.markdown(f"- 🔹 **Input:** {item['text']}")
-        st.markdown(f"   - ✨ **Translated:** {item['translated']}")
-        st.markdown(f"   - ❗ **Label:** {item['label']} with **{item['score']:.2%}** confidence")
-        st.markdown(f"   - 🔧 **Suggestion:** {item['reason']}")
-
-    # 雷达图
-    radar_df = pd.DataFrame({
-        "Category": ["Insult","Abuse","Discrimination","Hate Speech","Vulgarity"],
-        "Score": [0.7,0.4,0.3,0.5,0.6]
-    })
-    radar_fig = px.line_polar(radar_df, r='Score', theta='Category', line_close=True, title="⚠️ Risk Radar by Category")
-    radar_fig.update_traces(line_color='black')
-    st.plotly_chart(radar_fig)
-
-    # —— 新增：单词级冒犯性相关性分析 —— #
-    st.markdown("### 🧬 Word-level Offensive Correlation")
-
-    # 取最近一次翻译文本，按空格拆分单词
-    last_translated_text = st.session_state.history[-1]["translated"]
-    words = last_translated_text.split()
-
-    # 对每个单词进行分类并收集分数
-    word_scores = []
-    for word in words:
-        try:
-            res = classifier(word)[0]
-            word_scores.append({
-                "Word": word,
-                "Label": res["label"],
-                "Score": res["score"]
+            # 处理文本输入
+            text_input = text if text else ocr_text
+            analysis = classify_text_with_categories(text_input)
+            
+            # 更新历史记录
+            st.session_state.history.append({
+                "original": text_input,
+                "translated": analysis["translations"],
+                "overall": analysis["overall"],
+                "categories": analysis["categories"],
+                "word_analysis": analysis["word_analysis"]
             })
-        except Exception:
-            continue
+            
+            # 展示核心结果
+            st.markdown("**Main Prediction:**")
+            st.metric("Label", analysis["overall"]["label"], 
+                     delta=f"{analysis['overall']['score']:.2%} Confidence")
+            
+            # 新增：类别分布展示
+            st.markdown("### 📊 Offense Category Breakdown")
+            category_data = [{"Category": k, "Score": v} for k, v in analysis["categories"].items()]
+            fig = px.bar(category_data, x="Category", y="Score", 
+                        title="Category Contribution", 
+                        labels={"Score": "Probability"})
+            st.plotly_chart(fig)
 
-    if word_scores:
-        word_df = pd.DataFrame(word_scores)
-        word_df = word_df.sort_values(by="Score", ascending=False).reset_index(drop=True)
+        except Exception as e:
+            st.error(f"❌ Error: {str(e)}")
 
-        max_display = 5
-        # Streamlit 1.22+ 支持 st.toggle，若版本不支持可改用 checkbox
-        show_more = st.toggle("Show more words", value=False)
+# 优化后的雷达图生成
+if st.session_state.history:
+    # 聚合所有历史记录的类别数据
+    radar_data = {cat: [] for cat in OFFENSE_CATEGORIES}
+    for entry in st.session_state.history:
+        for cat, score in entry["categories"].items():
+            radar_data[cat].append(score)
+    
+    # 计算平均得分
+    avg_scores = {cat: sum(scores)/len(scores) if scores else 0 
+                 for cat, scores in radar_data.items()}
+    
+    # 构建雷达图
+    fig = px.line_polar(
+        pd.DataFrame(avg_scores, index=OFFENSE_CATEGORIES).reset_index(),
+        r='index', theta='OFFENSE_CATEGORIES',
+        line_close=True, title="📉 Offense Risk Radar Chart"
+    )
+    fig.update_traces(line_color='#FF4B4B')
+    st.plotly_chart(fig)
 
-        display_df = word_df if show_more else word_df.head(max_display)
-        # 隐藏边框并渲染 HTML 表格
-        st.markdown(
-            display_df.to_html(index=False, border=0),
-            unsafe_allow_html=True
-        )
+# 优化后的单词级分析
+if st.session_state.history:
+    # 聚合单词级数据
+    all_words = []
+    for entry in st.session_state.history:
+        all_words.extend(entry["word_analysis"])
+    
+    # 生成词云数据
+    word_counts = pd.DataFrame(all_words).groupby('word').agg({
+        'main_score': 'mean',
+        'offense_category': lambda x: x.mode()[0]
+    }).reset_index().sort_values('main_score', ascending=False)
+    
+    # 交互式词云展示
+    st.markdown("### 🧩 Offensive Word Analysis")
+    if not word_counts.empty:
+        top_words = word_counts.head(10)
+        fig = px.bar(top_words, x="word", y="main_score",
+                    color="offense_category",
+                    title="Top Offensive Words by Score")
+        st.plotly_chart(fig)
     else:
-        st.info("❕ No word-level analysis available.")
-else:
-    st.info("⚠️ No classification data available yet.")
+        st.info("No offensive words detected")