import streamlit as st
import pandas as pd
import plotly.express as px
import torch
import torch.nn as nn # Nạp thư viện PyTorch để định nghĩa lớp Custom
from transformers import AutoModel, AutoTokenizer
import os
import re
# Nạp module tiền xử lý độc lập từ file preprocessing.py của bạn
from preprocessing import preprocess_text_for_Visobert
# =========================================================================
# 1. CẤU HÌNH GIAO DIỆN STREAMLIT TOÀN MÀN HÌNH & CSS NỚI RỘNG SIDEBAR
# =========================================================================
st.set_page_config(page_title="Hệ thống Phát hiện Nội dung Độc hại", layout="wide")
st.markdown("""
""", unsafe_allow_html=True)
# =========================================================================
# 2. ĐỊNH NGHĨA LỚP TÙY CHỈNH ĐỒNG BỘ 100% VỚI CODE HUẤN LUYỆN CỦA BẠN
# =========================================================================
class ViSoBERTClassifier(nn.Module):
def __init__(self, visobert, num_labels=3):
super().__init__()
self.visobert = visobert
self.dropout = nn.Dropout(0.1)
self.fc = nn.Sequential(
nn.Linear(768, 256),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(256, num_labels)
)
def forward(self, input_ids, attention_mask):
outputs = self.visobert(
input_ids=input_ids,
attention_mask=attention_mask
)
cls = outputs.last_hidden_state[:, 0, :]
cls = self.dropout(cls)
return self.fc(cls)
@st.cache_resource
def load_sentiment_model():
model_path = "./best_model"
pt_file = os.path.join(model_path, "best_Visobert_mlp.pt")
if not os.path.exists(model_path) or not os.path.exists(pt_file):
st.error(f"❌ Không tìm thấy thư mục hoặc file trọng số PyTorch '{pt_file}'.")
st.stop()
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
visobert_backbone = AutoModel.from_pretrained("uitnlp/visobert")
model = ViSoBERTClassifier(visobert=visobert_backbone, num_labels=3)
# Nạp ma trận nơ-ron lớp MLP 3 nhãn chính xác lên CPU
state_dict = torch.load(pt_file, map_location=torch.device('cpu'))
model.load_state_dict(state_dict)
model.float()
model.eval()
id2label = {0: "🟢 CLEAN", 1: "🔴 TOXIC", 2: "🟠 OFFENSIVE"}
return tokenizer, model, id2label
tokenizer, model, id2label = load_sentiment_model()
# Hàm tính xác suất phần trăm khi người dùng gõ câu mới ở Sidebar
def predict_toxicity_probas(text):
text_processed = preprocess_text_for_Visobert(text)
# Đồng bộ khoảng trắng đầu vào tránh lệch nhãn
text_processed = text_processed.strip()
text_processed = re.sub(r"\s+", " ", text_processed)
if not text_processed.strip() or text_processed.lower() == 'nan':
return "🟢 CLEAN", {"CLEAN": 1.0, "TOXIC": 0.0, "OFFENSIVE": 0.0}
try:
inputs = tokenizer(text_processed, return_tensors="pt", padding=True, truncation=True, max_length=256)
with torch.no_grad():
logits = model(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
probs = torch.softmax(logits, dim=1).squeeze().tolist()
predicted_class_id = torch.argmax(logits, dim=1).item()
if isinstance(probs, float):
probs = [probs]
probas_dict = {
"CLEAN": float(probs[0]) if len(probs) > 0 else 0.0,
"TOXIC": float(probs[1]) if len(probs) > 1 else 0.0,
"OFFENSIVE": float(probs[2]) if len(probs) > 2 else 0.0
}
return id2label[predicted_class_id], probas_dict
except Exception:
return "🟢 CLEAN", {"CLEAN": 1.0, "TOXIC": 0.0, "OFFENSIVE": 0.0}
# =========================================================================
# 3. ĐỌC DỮ LIỆU ĐÃ GÁN NHÃN SẴN TỪ COLAB (TẬP TEST 3 CỘT: free_text, label_id, Pred)
# =========================================================================
def load_simulation_data_fixed():
file_path = "data_test_new.csv"
if os.path.exists(file_path):
try:
df = pd.read_csv(file_path, encoding='utf-8-sig')
df = df.dropna(subset=['free_text'])
df['free_text'] = df['free_text'].astype(str)
df = df[df['free_text'].str.lower() != 'nan']
df['comment'] = df['free_text'].apply(lambda x: re.sub(r'[\r\n\t]+', ' ', x))
def normalize_to_emoji_label(x):
x_str = str(x).strip().upper()
if x_str in ["0", "CLEAN"]: return "🟢 CLEAN"
if x_str in ["1", "TOXIC"]: return "🔴 TOXIC"
if x_str in ["2", "OFFENSIVE"]: return "🟠 OFFENSIVE"
return "🟢 CLEAN"
df['label'] = df['label_id'].apply(normalize_to_emoji_label)
df['Dự đoán'] = df['Pred'].apply(normalize_to_emoji_label)
df['is_manual'] = False
return df[['comment', 'label', 'Dự đoán', 'is_manual']]
except Exception:
pass
return pd.DataFrame(columns=['comment', 'label', 'Dự đoán', 'is_manual'])
# Khối lệnh nạp dữ liệu vào Session State
if "df_data" not in st.session_state:
with st.spinner("🔄 Hệ thống đang nạp tập dữ liệu..."):
st.session_state.df_data = load_simulation_data_fixed()
if "last_analysis" not in st.session_state:
st.session_state.last_analysis = None
if "show_sidebar_result" not in st.session_state:
st.session_state.show_sidebar_result = False
df_active = st.session_state.df_data
# =========================================================================
# 4. THIẾT KẾ GIAO DIỆN CHÍNH & SIDEBAR TỐI GIẢN CHUẨN ĐỊNH DẠNG MỚI
# =========================================================================
st.title("HỆ THỐNG PHÁT HIỆN NỘI DUNG ĐỘC HẠI")
st.caption("Đồ án môn học: Xử lý ngôn ngữ tự nhiên | Mô hình: VisoBERT| Đề tài: Phát hiện Toxic Comment")
# 🔥 ĐÃ SỬA: Ép phân cấp nút mẫu câu xếp dọc thông thoáng, sửa lỗi phân lớp ghi đè st.sidebar
st.sidebar.markdown("### Danh sách Mẫu Thử Nhanh")
sample_1 = "Hôm nay trời đẹp quá, chúc mọi người một ngày tốt lành!"
sample_2 = "thằng ad làm ăn như hạch, bực cả mình"
sample_3 = "xạo lol vừa thôi"
if st.sidebar.button(sample_1, use_container_width=True):
st.session_state.box_1_input = sample_1
st.session_state.box_2_input = sample_1
if st.sidebar.button(sample_2, use_container_width=True):
st.session_state.box_1_input = sample_2
st.session_state.box_2_input = sample_2
if st.sidebar.button(sample_3, use_container_width=True):
st.session_state.box_1_input = sample_3
st.session_state.box_2_input = sample_3
st.sidebar.markdown("---")
st.sidebar.markdown("### Demo phân loại toxic")
user_comment_1 = st.sidebar.text_area("Nhập nội dung kiểm tra (Lưu thẳng vào bảng):", height=80, key="box_1_input")
if st.sidebar.button("GỬI & XEM KẾT QUẢ", use_container_width=True):
if not user_comment_1.strip():
st.sidebar.error("⚠️ Vui lòng nhập nội dung!")
else:
pred_label, probas = predict_toxicity_probas(user_comment_1)
st.session_state.last_analysis = {"text": user_comment_1, "label": pred_label, "probas": probas}
st.session_state.show_sidebar_result = True
comment_cleaned = preprocess_text_for_Visobert(user_comment_1)
new_row = pd.DataFrame([{"comment": comment_cleaned, "label": "Phân loại (Demo)", "Dự đoán": pred_label, "is_manual": True}])
st.session_state.df_data = pd.concat([new_row, st.session_state.df_data], ignore_index=True)
st.rerun()
st.sidebar.markdown("---")
st.sidebar.markdown("### Demo tính năng Chặn độc hại")
user_comment_2 = st.sidebar.text_area("Nhập nội dung đăng tải (Mô phỏng bộ lọc bài):", height=80, key="box_2_input")
if st.sidebar.button("ĐĂNG BÌNH LUẬN", use_container_width=True):
if not user_comment_2.strip():
st.sidebar.error("⚠️ Vui lòng nhập nội dung!")
else:
pred_label, probas = predict_toxicity_probas(user_comment_2)
st.session_state.last_analysis = {"text": user_comment_2, "label": pred_label, "probas": probas}
st.session_state.show_sidebar_result = True
if "🔴 TOXIC" in pred_label or "🟠 OFFENSIVE" in pred_label:
st.sidebar.error("CHẶN BÀI: Nội dung bình luận vi phạm tiêu chuẩn cộng đồng! Hệ thống từ chối đăng tải. Vui lòng nhập lại câu từ lịch sự hơn.")
else:
comment_cleaned = preprocess_text_for_Visobert(user_comment_2)
new_row = pd.DataFrame([{"comment": comment_cleaned, "label": "Chặn độc hại (Demo)", "Dự đoán": pred_label, "is_manual": True}])
st.session_state.df_data = pd.concat([new_row, st.session_state.df_data], ignore_index=True)
st.sidebar.success("DUYỆT ĐĂNG: Bình luận an toàn! Đã cập nhật lên hệ thống.")
st.rerun()
# Khối hiển thị kết quả phân tích bằng hàm progress (Có màu sắc theo nhãn cảm xúc)
if st.session_state.show_sidebar_result and st.session_state.last_analysis:
res = st.session_state.last_analysis
lbl_text = res["label"].replace("🟢 ", "").replace("🔴 ", "").replace("🟠 ", "")
max_proba = float(res["probas"][lbl_text]) * 100
st.sidebar.markdown("---")
st.sidebar.markdown(f"#### KẾT QUẢ PHÂN TÍCH")
st.sidebar.markdown(f"**Trạng thái:** {res['label']}")
st.sidebar.markdown(f"**Độ tin cậy:** `{max_proba:.2f}%`")
st.sidebar.markdown("
", unsafe_allow_html=True)
st.sidebar.markdown(f"CLEAN: {res['probas']['CLEAN']*100:.1f}%", unsafe_allow_html=True)
st.sidebar.markdown(f"""
""", unsafe_allow_html=True)
st.sidebar.markdown(f"TOXIC: {res['probas']['TOXIC']*100:.1f}%", unsafe_allow_html=True)
st.sidebar.markdown(f"""
""", unsafe_allow_html=True)
st.sidebar.markdown(f"OFFENSIVE: {res['probas']['OFFENSIVE']*100:.1f}%", unsafe_allow_html=True)
st.sidebar.markdown(f"""
""", unsafe_allow_html=True)
# --- KHỐI KPI THỐNG KÊ TỔNG HỢP ---
total_cnt = len(df_active)
clean_cnt = len(df_active[df_active["Dự đoán"] == "🟢 CLEAN"])
toxic_cnt = len(df_active[df_active["Dự đoán"] == "🔴 TOXIC"])
off_cnt = len(df_active[df_active["Dự đoán"] == "🟠 OFFENSIVE"])
m1, m2, m3, m4 = st.columns(4)
m1.metric("Tổng số bình luận", f"{total_cnt:,} câu")
m2.metric("Số lượng CLEAN", f"{clean_cnt:,} câu", f"{(clean_cnt/total_cnt)*100:.1f}%" if total_cnt else "0%")
m3.metric("Số lượng TOXIC", f"{toxic_cnt:,} câu", f"{(toxic_cnt/total_cnt)*100:.1f}%" if total_cnt else "0%", delta_color="inverse")
m4.metric("Số lượng OFFENSIVE", f"{off_cnt:,} câu", f"{(off_cnt/total_cnt)*100:.1f}%" if total_cnt else "0%")
# =========================================================================
# BẢNG 1: CẢNH BÁO KHẨN CẤP
# =========================================================================
st.markdown("---")
df_emergency = df_active[(df_active["is_manual"] == True) & (df_active["Dự đoán"].isin(["🔴 TOXIC", "🟠 OFFENSIVE"]))]
st.subheader(f"⚠️ CẢNH BẢO NỘI DUNG ĐỘC HẠI: {len(df_emergency)}")
if not df_emergency.empty:
df_disp_emergency = df_emergency.copy()[["comment", "Dự đoán"]]
else:
df_disp_emergency = pd.DataFrame(columns=["comment", "Dự đoán"])
st.dataframe(
df_disp_emergency, use_container_width=True, hide_index=True,
column_config={
"comment": st.column_config.TextColumn("Comment", width="large"),
"Dự đoán": st.column_config.TextColumn("nhãn dự đoán", width="small")
}
)
# =========================================================================
# KHỐI BIỂU ĐỒ TRỰC QUAN
# =========================================================================
st.markdown("---")
st.subheader("Biểu đồ kết quả phân loại")
if not df_active.empty:
color_map_global = {"🟢 CLEAN": "#2ecc71", "🔴 TOXIC": "#e74c3c", "🟠 OFFENSIVE": "#e67e22"}
sentiment_counts = df_active["Dự đoán"].value_counts()
sentiment_df = pd.DataFrame({"Phân lớp": sentiment_counts.index, "Solo": sentiment_counts.values})
fig_bar = px.bar(sentiment_df, x="Phân lớp", y="Solo", color="Phân lớp", color_discrete_map=color_map_global)
fig_bar.update_layout(xaxis=dict(title=dict(text="Nhãn phân loại", standoff=45), automargin=True), yaxis=dict(title_text="Tổng số dòng"), margin=dict(b=100), legend_title_text="Phân lớp")
st.plotly_chart(fig_bar, use_container_width=True)
# =========================================================================
# BẢNG 2: BẢNG TỔNG HỢP BÌNH LUẬN VÀ PHÂN LOẠI
# =========================================================================
st.markdown("---")
st.subheader("Bảng tổng hợp bình luận và phân loại")
if not df_active.empty:
df_display_all = df_active.copy()[["comment", "label", "Dự đoán", "is_manual"]]
else:
df_display_all = pd.DataFrame(columns=["comment", "label", "Dự đoán", "is_manual"])
st.dataframe(
df_display_all, use_container_width=True, hide_index=True,
column_config={
"comment": st.column_config.TextColumn("Comment", width="large"),
"label": st.column_config.TextColumn("Nhãn gốc thực tế", width="small"),
"Dự đoán": st.column_config.TextColumn("nhãn dự đoán", width="small"),
"is_manual": st.column_config.CheckboxColumn("is_manual", width=80)
}
)