Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,79 +1,85 @@
 import os
 import streamlit as st
-from transformers import DistilBertTokenizer, TFDistilBertModel
-from tensorflow.keras.models import load_model
 import numpy as np
 import tensorflow as tf
+from tensorflow.keras.layers import Input, Lambda, Dense
+from tensorflow.keras.models import Model
+from transformers import DistilBertTokenizer, TFDistilBertModel
 
 # =======================
-# ĐẶT LẠI CACHE ĐỂ TRÁNH LỖI TRÊN SPACES
+# CẤU HÌNH
 # =======================
-os.environ['TRANSFORMERS_CACHE'] = './cache'
+MAX_LEN = 400
+WEIGHTS_PATH = "src/model_Adam.h5"
+TOKENIZER_PATH = "src"
+CACHE_DIR = "./cache"
 
 # =======================
-# CẤU HÌNH
+# TRÁNH LỖI GHI CACHE
 # =======================
-MAX_LEN = 400
-MODEL_PATH = "src/model_Adam.h5"
-TOKENIZER_PATH = "src/"  
+os.environ['TRANSFORMERS_CACHE'] = CACHE_DIR
+
 # =======================
-# LOAD TOKENIZER & TRANSFORMER
+# TẢI TOKENIZER
 # =======================
 @st.cache_resource
 def load_tokenizer():
     return DistilBertTokenizer.from_pretrained(TOKENIZER_PATH)
 
-@st.cache_resource
-def load_transformer():
-    return TFDistilBertModel.from_pretrained(TOKENIZER_PATH)
-
-# =======================
-# ĐỊNH NGHĨA transformer_layer
-# =======================
-def transformer_layer(inputs):
-    input_ids, attention_mask = inputs
-    transformer = load_transformer()
-    outputs = transformer(input_ids=input_ids, attention_mask=attention_mask)
-    return outputs.last_hidden_state[:, 0, :]
+tokenizer = load_tokenizer()
 
 # =======================
-# LOAD MÔ HÌNH PHÂN LOẠI CẢM XÚC
+# TẠO MÔ HÌNH (PHẢI GIỐNG KHI TRAIN)
 # =======================
 @st.cache_resource
-def load_sentiment_model():
-    return load_model(MODEL_PATH, custom_objects={'transformer_layer': transformer_layer})
+def create_model_and_load_weights():
+    transformer = TFDistilBertModel.from_pretrained("distilbert-base-uncased", cache_dir=CACHE_DIR)
+
+    input_ids = Input(shape=(MAX_LEN,), dtype=tf.int32, name="input_ids")
+    attention_mask = Input(shape=(MAX_LEN,), dtype=tf.int32, name="attention_mask")
+
+    def transformer_layer(inputs):
+        ids, mask = inputs
+        outputs = transformer(input_ids=ids, attention_mask=mask)
+        return outputs.last_hidden_state[:, 0, :]  # Lấy CLS token
+
+    cls_output = Lambda(transformer_layer)([input_ids, attention_mask])
+    output = Dense(1, activation='sigmoid')(cls_output)
+
+    model = Model(inputs=[input_ids, attention_mask], outputs=output)
+    model.load_weights(WEIGHTS_PATH)
+    return model
+
+model = create_model_and_load_weights()
 
 # =======================
-# TIỀN XỬ LÝ VĂN BẢN
+# TIỀN XỬ LÝ
 # =======================
-tokenizer = load_tokenizer()
-
 def preprocess(text):
     tokens = tokenizer(
         text,
         max_length=MAX_LEN,
-        padding='max_length',
+        padding="max_length",
         truncation=True,
-        return_tensors='tf'
+        return_tensors="tf"
     )
     return {
-        'input_ids': tokens['input_ids'],
-        'attention_mask': tokens['attention_mask']
+        "input_ids": tokens["input_ids"],
+        "attention_mask": tokens["attention_mask"]
     }
 
 # =======================
-# STREAMLIT APP
+# GIAO DIỆN STREAMLIT
 # =======================
-st.title("🎬 Sentiment Analysis — Đánh giá phim")
+st.title("🎬 Phân tích cảm xúc đánh giá phim")
 
-user_input = st.text_area("Nhập nội dung đánh giá phim của bạn:", height=150)
+user_input = st.text_area("Nhập đánh giá phim của bạn:", height=150)
 
 if st.button("Dự đoán cảm xúc"):
     if not user_input.strip():
-        st.warning("Vui lòng nhập nội dung trước khi dự đoán.")
+        st.warning("Vui lòng nhập nội dung.")
     else:
         with st.spinner("Đang xử lý..."):
-            model = load_sentiment_model()
             inputs = preprocess(user_input)
             prob = model.predict(inputs)[0][0]
             label = "TÍCH CỰC 😊" if prob >= 0.5 else "TIÊU CỰC 😞"