Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,113 +1,96 @@
 # app.py
 import streamlit as st
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# --- Dữ liệu Demo (Chuẩn bị sẵn) ---
-# Dữ liệu này được lấy sẵn để demo nhanh và ổn định.
-DEMO_DATA = {
-    "Ví dụ 1: Dấu hiệu trầm cảm rõ ràng": {
-        "post": "Lately, I just feel so empty and numb inside. Nothing brings me joy anymore, not even the things I used to love. It's a struggle just to get out of bed.",
-        "mentallama_output": {
-            "answer": "Depression",
-            "reasoning": "The post conveys a deep sense of emotional numbness, loss of interest (anhedonia), and lack of motivation, which are core symptoms of depression."
-        },
-        "mentalroberta_output": "Depression"
-    },
-    "Ví dụ 2: Dấu hiệu stress tinh tế": {
-        "post": "Another sleepless night. My mind just won't shut off. I feel so tired but I can't rest. Work is piling up and I just don't have the energy to face it.",
-        "mentallama_output": {
-            "answer": "Stress",
-            "reasoning": "The user mentions sleepless nights, a racing mind, and feeling overwhelmed by work. These are classic indicators of high stress levels."
-        },
-        "mentalroberta_output": "Stress"
-    },
-    "Ví dụ 3: Bài đăng tích cực, không có vấn đề": {
-        "post": "Just finished a great hike with friends! The view from the top was incredible. Feeling so refreshed and happy.",
-        "mentallama_output": {
-            "answer": "None",
-            "reasoning": "The post expresses positive emotions, mentions social activity (hiking with friends), and feelings of refreshment. There are no indicators of mental distress."
-        },
-        "mentalroberta_output": "None"
-    },
-    "Ví dụ 4: Tự tử (Nghiêm trọng)": {
-        "post": "I can't take this pain anymore. I've been thinking about ending it all. It feels like the only way out.",
-        "mentallama_output": {
-            "answer": "Suicide",
-            "reasoning": "The post contains explicit suicidal ideation ('thinking about ending it all') and expresses feelings of unbearable pain and hopelessness, indicating a high-risk situation."
-        },
-        "mentalroberta_output": "Suicide"
+# --- Tối ưu hóa: Tải mô hình chỉ một lần duy nhất ---
+# @st.cache_resource là một "bảo bối" của Streamlit.
+# Nó đảm bảo hàm này chỉ chạy một lần khi ứng dụng khởi động.
+# Lần sau khi người dùng tương tác, mô hình đã có sẵn trong bộ nhớ.
+@st.cache_resource
+def load_model():
+    model_id = "SteveKGYang/MentaLLaMA-chat-7B"
+    
+    # Cấu hình Quantization 4-bit để giảm RAM
+    bnb_config = {
+        "load_in_4bit": True,
+        "bnb_4bit_quant_type": "nf4",
+        "bnb_4bit_compute_dtype": torch.bfloat16,
+        "bnb_4bit_use_double_quant": False,
     }
-}
-
-# --- Giao diện ứng dụng Streamlit ---
-st.set_page_config(layout="wide", page_title="MentaLLaMA Demo")
 
-# Header
-st.title("🔬 Demo: MentaLLaMA - Phân Tích Sức Khỏe Tâm Thần Có Thể Giải Thích")
-st.markdown("""
-Chào mừng đến với bản demo của **MentaLLaMA**!
+    # Tải mô hình với cấu hình đã thiết lập
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        quantization_config=bnb_config,
+        device_map="auto", # Tự động phân bổ lên các thiết bị có sẵn (CPU)
+        trust_remote_code=True
+    )
 
-Demo này minh họa sự khác biệt cốt lõi giữa:
-1.  **Phương pháp truyền thống "Hộp đen"** (như MentalRoBERTa): Chỉ đưa ra dự đoán.
-2.  **MentaLLaMA**: Đưa ra cả dự đoán và **lời giải thích chi tiết**.
+    # Tải tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-*Lưu ý: Demo này sử dụng kết quả được tạo sẵn để đảm bảo tốc độ và sự ổn định.*
-""")
+    return model, tokenizer
 
-# --- Phần tương tác ---
-st.sidebar.header("Bảng điều khiển")
-option = st.sidebar.selectbox(
-    'Chọn một bài đăng mẫu để phân tích:',
-    list(DEMO_DATA.keys()),
-    index=0
-)
+# --- Hàm để chạy suy luận ---
+def run_inference(post_text, model, tokenizer):
+    # Đây là cấu trúc prompt mà MentaLLaMA được huấn luyện để tuân theo.
+    # Việc tuân thủ đúng prompt format là RẤT QUAN TRỌNG.
+    prompt = f"""### Human: Analyze the following post and provide a diagnosis along with a detailed reasoning. Post: {post_text} ### Assistant:"""
 
-# Hiển thị bài đăng được chọn
-selected_example = DEMO_DATA[option]
-post_text = selected_example["post"]
+    # Mã hóa prompt thành các token
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
 
-st.subheader("📝 Nội dung bài đăng được chọn:")
-st.info(post_text)
+    # Chạy mô hình để tạo ra kết quả
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=256,  # Giới hạn độ dài của câu trả lời
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.eos_token_id
+    )
 
-# Nút phân tích
-if st.sidebar.button('Chạy Phân Tích So Sánh', type="primary"):
-    st.header("📊 Kết quả phân tích so sánh")
+    # Giải mã kết quả từ token về lại văn bản
+    result_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    col1, col2 = st.columns(2, gap="large")
+    # Xử lý chuỗi để chỉ lấy phần trả lời của Assistant
+    # Kết quả thô sẽ có dạng: "### Human: ... ### Assistant: Answer: ... Reasoning: ..."
+    # Chúng ta cần cắt bỏ phần prompt đi.
+    try:
+        assistant_response = result_text.split("### Assistant:")[1].strip()
+        return assistant_response
+    except IndexError:
+        return "Lỗi: Không thể phân tích cú pháp đầu ra của mô hình."
 
-    # Cột 1: Phương pháp truyền thống (Hộp đen)
-    with col1:
-        st.subheader("Mô hình 'Hộp đen' (ví dụ: MentalRoBERTa)")
-        
-        st.metric(label="Dự đoán", value=selected_example['mentalroberta_output'])
-        
-        with st.expander("Phân tích về phương pháp này", expanded=True):
-            st.warning("⚠️ **Không có lời giải thích!**")
-            st.write("""
-            Mô hình chỉ đưa ra kết quả cuối cùng. Chúng ta không biết **tại sao** nó lại đưa ra dự đoán này. 
-            - **Khó tin tưởng:** Làm sao để chắc chắn nó không mắc sai lầm?
-            - **Khó xác thực:** Bác sĩ không có cơ sở để kiểm tra lại logic của AI.
-            - **Rủi ro cao:** Trong lĩnh vực y tế, một quyết định không thể giải thích là rất nguy hiểm.
-            """)
+# --- Giao diện ứng dụng Streamlit ---
+st.set_page_config(layout="wide", page_title="MentaLLaMA Live Demo")
 
-    # Cột 2: MentaLLaMA (Giải thích được)
-    with col2:
-        st.subheader("MentaLLaMA (Mô hình giải thích được)")
-        
-        st.metric(label="Dự đoán", value=selected_example['mentallama_output']['answer'])
-        
-        st.write("**Lý do giải thích (Reasoning):**")
-        st.success(f"“{selected_example['mentallama_output']['reasoning']}”")
+st.title("🔬 MentaLLaMA Live Demo - Chạy mô hình thật")
+st.markdown("""
+Chào mừng đến với bản demo **chạy thật** của **MentaLLaMA**! 
+- **Nhập** một đoạn văn bản (bằng tiếng Anh) vào ô bên dưới.
+- **Nhấn nút** để mô hình MentaLLaMA-7B phân tích trực tiếp.
+- **Lưu ý:** Vì chạy trên CPU miễn phí, quá trình suy luận có thể mất **1-2 phút**. Vui lòng kiên nhẫn!
+""")
+
+# Tải mô hình (sẽ hiển thị thanh tiến trình lần đầu)
+with st.spinner("Đang tải mô hình MentaLLaMA-7B (lần đầu có thể mất vài phút)..."):
+    model, tokenizer = load_model()
+
+st.success("Mô hình đã sẵn sàng!")
+
+# Ô nhập liệu cho người dùng
+user_input = st.text_area("Nhập bài đăng của bạn vào đây (tiếng Anh):", "Lately, I just feel so empty and numb inside. Nothing brings me joy anymore...", height=150)
+
+# Nút chạy suy luận
+if st.button('Chạy Phân Tích MentaLLaMA', type="primary"):
+    if user_input:
+        with st.spinner("MentaLLaMA đang suy luận, quá trình này có thể mất 1-2 phút..."):
+            result = run_inference(user_input, model, tokenizer)
         
-        with st.expander("Phân tích về phương pháp này", expanded=True):
-            st.info("💡 **Cung cấp lý do chi tiết!**")
-            st.write("""
-            MentaLLaMA không chỉ dự đoán mà còn chỉ ra các **bằng chứng cụ thể** trong văn bản.
-            - **Tăng sự tin tưởng:** Lời giải thích hợp lý giúp chúng ta tin vào kết quả.
-            - **Hỗ trợ chuyên gia:** Bác sĩ có thể nhanh chóng xác thực logic và đưa ra quyết định cuối cùng.
-            - **An toàn hơn:** Sự minh bạch giúp giảm thiểu rủi ro và tăng tính trách nhiệm.
-            """)
-else:
-    st.info("Hãy chọn một ví dụ và nhấn nút 'Chạy Phân Tích So Sánh' ở thanh bên trái.")
+        st.subheader("Kết quả phân tích từ MentaLLaMA:")
+        st.success(result)
+    else:
+        st.warning("Vui lòng nhập nội dung bài đăng.")
 
 # Footer
 st.markdown("---")