Update app.py

app.py

@@ -2,56 +2,72 @@ import streamlit as st
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
-# 设置设备（优先用GPU）
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-pipe_device = 0 if torch.cuda.is_available() else -1
-
-# 加载自定义情感分析模型
-sentiment_pipeline = pipeline(
-    "text-classification",
-    model="KeviniveK/CustomModel_IMDB",
-    device=pipe_device
-)
-
-# 加载 NLLB 翻译模型：facebook/nllb-200-distilled-600M
-translation_model_name = "facebook/nllb-200-distilled-600M"
-translation_tokenizer = AutoTokenizer.from_pretrained(translation_model_name)
-translation_model = AutoModelForSeq2SeqLM.from_pretrained(translation_model_name).to(device)
-
-# 翻译函数（英文 → 中文简体）
-def translate_en_to_zh(text):
-    inputs = translation_tokenizer(
+# ------------ 缓存模型加载 ------------
+@st.cache_resource(show_spinner=False)
+def load_sentiment_pipeline():
+    device = 0 if torch.cuda.is_available() else -1
+    return pipeline(
+        "text-classification",
+        model="KeviniveK/CustomModel_IMDB",
+        device=device
+    )
+
+@st.cache_resource(show_spinner=False)
+def load_translation_model():
+    model_name = "facebook/nllb-200-distilled-600M"
+    # 加载 tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    # 加载 model，并让 accelerate 自动把各层放到合适设备上
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        model_name,
+        device_map="auto",           # 自动分配设备
+        torch_dtype=torch.float16    # 半精度减少显存占用（可选）
+    )
+    return tokenizer, model
+
+# ------------ 翻译函数 ------------
+@st.cache_data(show_spinner=False)
+def translate_en_to_zh(text: str, tokenizer, model):
+    # 直接将 inputs 放到 model.device 上
+    inputs = tokenizer(
         text,
         return_tensors="pt",
         truncation=True,
         padding=True
-    ).to(device)
+    )
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
-    translated_tokens = translation_model.generate(
+    translated_tokens = model.generate(
         **inputs,
-        forced_bos_token_id=translation_tokenizer.lang_code_to_id["zho_Hans"]
+        forced_bos_token_id=tokenizer.lang_code_to_id["zho_Hans"]
     )
+    return tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
 
-    translated_text = translation_tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
-    return translated_text
-
-# Streamlit 主逻辑
+# ------------ Streamlit App 主逻辑 ------------
 def main():
     st.set_page_config(page_title="IMDB 影评分析与翻译", layout="centered")
     st.title("🎬 英文影评情感分析 + 中文翻译")
     st.write("请输入英文影评内容，系统将输出情感分析结果与对应的中文翻译。")
 
-    user_input = st.text_area("英文影评输入", height=150, placeholder="e.g. The story was powerful and the acting was phenomenal.")
+    user_input = st.text_area(
+        "英文影评输入",
+        height=150,
+        placeholder="e.g. The story was powerful and the acting was phenomenal."
+    )
 
     if user_input:
         with st.spinner("分析处理中..."):
-            # 情感分析
+            # 加载已经缓存的模型
+            sentiment_pipeline = load_sentiment_pipeline()
+            tokenizer, translation_model = load_translation_model()
+
+            # 1. 情感分析
             sentiment_result = sentiment_pipeline(user_input)[0]
             sentiment = sentiment_result["label"]
             confidence = sentiment_result["score"]
 
-            # 翻译
-            translation = translate_en_to_zh(user_input)
+            # 2. 翻译
+            translation = translate_en_to_zh(user_input, tokenizer, translation_model)
 
         st.success("分析完成 ✅")
         st.markdown(f"**情感分类**：{sentiment}")
@@ -60,4 +76,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-