Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,106 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import os
+# Thiết lập env trước khi import bất kỳ module nào dùng Streamlit hoặc Transformers
+os.environ['TRANSFORMERS_CACHE'] = '/cache/hf_cache'
+os.environ['HF_HOME'] = '/cache/hf_cache'
+os.environ['XDG_CACHE_HOME'] = '/cache/.cache'
+os.environ['STREAMLIT_CONFIG_DIR'] = '/cache/.streamlit'
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+import asyncio
+import nest_asyncio
+nest_asyncio.apply()
+
+import tempfile
 import streamlit as st
+import librosa
+import torch
+import pandas as pd
+from transformers import (
+    Wav2Vec2Processor, Wav2Vec2ForCTC,
+    WhisperProcessor, WhisperForConditionalGeneration,
+    AutoTokenizer, AutoModelForTokenClassification,
+    AutoProcessor, AutoModelForSpeechSeq2Seq,
+    pipeline,
+)
+
+# disable torch dynamo for stability
+import torch._dynamo
+torch._dynamo.disable()
+
+# --- Configuration: model paths ---
+ASR_MODELS = {
+    "PhoWhisper": "Huydb/phowhisper-toxic",  # ensure this repo has processor_config.json
+}
+TSD_MODELS = {
+    "PhoBERT": "Huydb/PhoBERT-toxic",
+}
+
+# --- Load ASR processors & models (cached) ---
+@st.cache_resource
+def load_asr(path):
+    proc = WhisperProcessor.from_pretrained(path, cache_dir=os.environ['HF_HOME'])
+    mod = WhisperForConditionalGeneration.from_pretrained(path, cache_dir=os.environ['HF_HOME'])
+    return proc, mod
+
+
+asr_path = "Huydb/phowhisper-toxic"
+asr_processor, asr_model = load_asr(asr_path)
+
+# --- Load TSD tokenizers & models ---
+@st.cache_resource
+def load_tsd(path):
+    tok = AutoTokenizer.from_pretrained(path, cache_dir=os.environ['HF_HOME'])
+    mod = AutoModelForTokenClassification.from_pretrained(path, num_labels=2, cache_dir=os.environ['HF_HOME'])
+    return tok, mod
+
+tsd_path = "Huydb/PhoBERT-toxic"
+tsd_tokenizer, tsd_model = load_tsd(tsd_path)
+
+# --- Streamlit UI ---
+st.markdown("""
+<style> /* CSS animation & button */
+@keyframes bgfade {0%{background-color:white;}50%{background-color:#889ECE;}100%{background-color:white;}}
+html, body, .reportview-container, .main {height:100%!important; margin:0; padding:0; animation:bgfade 10s ease infinite;}
+ div.stButton>button:first-child{background-color:red!important;color:white!important;border:none;}
+</style>
+""", unsafe_allow_html=True)
+
+st.title("🔊🤬 Toxic Spans Detection from Audio")
+uploaded_audio = st.file_uploader("1. Upload a WAV audio file", type=["wav"])
+if not uploaded_audio:
+    st.info("Please upload a WAV audio file to begin.")
+    st.stop()
+
+with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tfile:
+    tfile.write(uploaded_audio.read())
+    audio_path = tfile.name
+st.success("Audio uploaded.")
+st.audio(audio_path, format='audio/wav')
+
+
+# Process button
+def highlight_toxic_span(words, labels):
+    sen_hide = ""
+    for word, label in zip(words, labels):
+        if label == 1:
+            sen_hide += "*"*len(word) + " "
+        else:
+            sen_hide += word + " "
+    return sen_hide.strip()
+
+if st.button("Transcript and Detect Toxic Spans Now"):
+    waveform, _ = librosa.load(audio_path, sr=16000)
+    input_features = proc(waveform, return_tensors="pt", sampling_rate=16000).input_features.to("cpu")
+    predicted_ids = mod.generate(input_features)
+    transcript_text = proc.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+
+    st.subheader("Result")
+    enc = tsd_tokenizer(list([transcript_text]), is_split_into_words=True,
+              padding='max_length', truncation=True,
+              max_length=len(list(transcript_text)), return_tensors="pt")
+    with torch.no_grad(): 
+        logits = tsd_model(input_ids=enc.input_ids, attention_mask=enc.attention_mask).logits
+    labels = logits.argmax(-1)[0].cpu().tolist()
+    sen_hide = highlight_toxic_span(transcript_text.split(), labels)
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+    st.markdown(f"<h1 style='text-align: center; color: red;'>{sen_hide}</h1>", unsafe_allow_html=True)