Update app.py

app.py

@@ -1,124 +1,219 @@
-# 新增：预定义冒犯性类别映射（根据雷达图需求）
-OFFENSE_CATEGORIES = {
-    "Insult": ["侮辱", "贬低", "讽刺"],
-    "Abuse": ["辱骂", "攻击性语言", "脏话"],
-    "Discrimination": ["歧视性言论", "种族歧视", "性别歧视"],
-    "Hate Speech": ["仇恨言论", "暴力煽动"],
-    "Vulgarity": ["低俗用语", "色情暗示"]
+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
+
+# ✅ 新增维度定义
+OFFENSIVE_CATEGORIES = {
+    "Insult": ["蠢货", "白痴", "废物"],
+    "Abuse": ["去死", "打死", "宰了你"],
+    "Discrimination": ["女司机", "娘娘腔", "黑鬼"],
+    "HateSpeech": ["灭族", "屠杀", "灭绝"],
+    "Vulgarity": ["艹", "sb", "尼玛"]
 }
 
-# 修改分类函数以支持多维度分析
-def classify_text_with_categories(text: str):
-    results = classifier(text)
-    category_scores = {category: 0 for category in OFFENSE_CATEGORIES}
+# ✅ 模型初始化（保持原有结构）
+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()
+
+model_options = {
+    "Toxic-BERT": "unitary/toxic-bert",
+    "Roberta Offensive": "cardiffnlp/twitter-roberta-base-offensive",
+    "BERT Emotion": "bhadresh-savani/bert-base-go-emotion"
+}
+
+# ✅ 动态评分算法
+def dynamic_scoring(text: str, classifier):
+    scores = {k: 0.0 for k in OFFENSIVE_CATEGORIES}
+    
+    for category, keywords in OFFENSIVE_CATEGORIES.items():
+        for kw in keywords:
+            if kw in text:
+                scores[category] += 0.3
     
-    # 多维度评分
-    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():
+    words = text.split()
+    for word in words:
         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
+            if res["label"] in scores:
+                scores[res["label"]] += res["score"] * 0.7
+        except: pass
+    
+    max_score = max(scores.values()) or 1
+    return {k: round(v/max_score, 2) for k,v in scores.items()}
 
-    return {
-        "translations": text,
-        "overall": results[0],
-        "categories": category_scores,
-        "word_analysis": word_analysis
-    }
+# ✅ 分类函数改造
+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()
 
-# 优化后的分类处理
-if st.button("🚦 Analyze Text"):
-    with st.spinner("🔍 Processing..."):
-        try:
-            # 处理文本输入
-            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"]
-            })
-            
-            # 展示核心结果
-            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)
-
-        except Exception as e:
-            st.error(f"❌ Error: {str(e)}")
-
-# 优化后的雷达图生成
-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)
+    result = classifier(translated_text)[0]
+    label = result["label"]
+    score = result["score"]
+    reasoning = f"The sentence was flagged as '{label}' due to potentially offensive phrases."
     
-    # 计算平均得分
-    avg_scores = {cat: sum(scores)/len(scores) if scores else 0 
-                 for cat, scores in radar_data.items()}
+    # 新增维度分析
+    category_scores = dynamic_scoring(translated_text, classifier)
+    
+    st.session_state.history.append({
+        "text": text,
+        "translated": translated_text,
+        "label": label,
+        "score": score,
+        "reason": reasoning,
+        "scores": category_scores
+    })
+    return translated_text, label, score, reasoning, category_scores
+
+# ✅ 可视化生成函数
+def generate_radar_chart(scores_dict: dict):
+    radar_df = pd.DataFrame({
+        "Category": list(scores_dict.keys()),
+        "Score": list(scores_dict.values())
+    })
     
-    # 构建雷达图
     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"
+        radar_df, 
+        r='Score', 
+        theta='Category',
+        line_close=True,
+        color_discrete_sequence=['#FF6B6B'],
+        title="🛡️ Multi-Dimensional Offensive Analysis"
     )
-    fig.update_traces(line_color='#FF4B4B')
-    st.plotly_chart(fig)
+    fig.update_layout(
+        polar=dict(
+            radialaxis=dict(
+                visible=True,
+                range=[0, 1],
+                tickvals=[0, 0.3, 0.7, 1],
+                ticktext=["Safe", "Caution", "Risk", "Danger"]
+            )),
+        showlegend=False
+    )
+    return fig
+
+# ✅ 页面配置（保持原有结构）
+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 = []
+
+# 主页面逻辑
+st.title("🚨 Emoji Offensive Text Detector & Analysis Dashboard")
 
-# 优化后的单词级分析
+# 文本输入
+st.subheader("1. 输入与分类")
+default_text = "你是🐷"
+text = st.text_area("Enter sentence with emojis:", value=default_text, height=150)
+
+if st.button("🚦 Analyze Text"):
+    with st.spinner("🔍 Processing..."):
+        try:
+            translated, label, score, reason, category_scores = classify_emoji_text(text)
+            # 展示基础结果
+            st.markdown("**Translated sentence:**")
+            st.code(translated, language="text")
+            # 展示雷达图
+            st.plotly_chart(generate_radar_chart(category_scores))
+
+# 图片上传与 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:
-    # 聚合单词级数据
-    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)
+    # 展示历史记录
+    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"]
+            })
+        except Exception:
+            continue
+
+    if word_scores:
+        word_df = pd.DataFrame(word_scores)
+        word_df = word_df.sort_values(by="Score", ascending=False).reset_index(drop=True)
+
+        max_display = 5
+        # Streamlit 1.22+ 支持 st.toggle，若版本不支持可改用 checkbox
+        show_more = st.toggle("Show more words", value=False)
+
+        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
+        )
     else:
-        st.info("No offensive words detected")
+        st.info("❕ No word-level analysis available.")
+else:
+    st.info("⚠️ No classification data available yet.")