Update app.py

app.py

@@ -1,58 +1,125 @@
 import gradio as gr
 import numpy as np
-from sentence_transformers import SentenceTransformer
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
 import os
+import torch
+import pandas as pd
+from sentence_transformers import SentenceTransformer
 from huggingface_hub import hf_hub_download
 import faiss
+import time
 
-# 设置内存交换参数
+# 减少内存占用
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:32"
+torch.set_num_threads(1)
+
+# 初始化空模型
+model = None
+index = None
+metadata = None
+tokenizer = None
+retriever = None
+
+def load_resources():
+    """按需加载资源"""
+    global model, index, metadata, tokenizer, retriever
+    
+    # 仅当需要时加载
+    if model is None:
+        print("正在加载句子嵌入模型...")
+        token = os.getenv("HF_TOKEN")
+        model = SentenceTransformer("all-MiniLM-L6-v2", use_auth_token=token)
+        print("句子模型加载完成")
 
-# 设置默认线程数
-torch.set_num_threads(1)  # 减少内存占用
+    if index is None:
+        print("正在下载FAISS索引...")
+        INDEX_PATH = hf_hub_download(
+            repo_id="GOGO198/GOGO_rag_index",
+            filename="faiss_index.bin",
+            cache_dir="/data",
+            use_auth_token=os.getenv("HF_TOKEN")
+        )
+        index = faiss.read_index(INDEX_PATH)
+        print("FAISS索引加载完成")
 
-# 使用更轻量的模型 - 添加 token 参数
-model_name = "all-MiniLM-L6-v2"
-token = os.getenv("HF_TOKEN")  # 从环境变量获取令牌
-model = SentenceTransformer(model_name, use_auth_token=token) if token else None
+    if metadata is None:
+        print("正在下载元数据...")
+        METADATA_PATH = hf_hub_download(
+            repo_id="GOGO198/GOGO_rag_index",
+            filename="metadata.csv",
+            cache_dir="/data",
+            use_auth_token=os.getenv("HF_TOKEN")
+        )
+        metadata = pd.read_csv(METADATA_PATH)
+        print("元数据加载完成")
+
+    # 延迟加载RAG组件
+    if tokenizer is None:
+        tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
+        
+    if retriever is None:
+        retriever = RagRetriever.from_pretrained(
+            "facebook/rag-sequence-nq",
+            index_name="custom",
+            index_paths=["/data/rag_index.faiss"]  # 使用已加载的索引
+        )
 
 def predict(vector):
-    # 加载本地索引
-    retriever = RagRetriever.from_pretrained(
-        "facebook/rag-sequence-nq",
-        index_name="custom",
-        index_paths=["rag_index.faiss"]
-    )
-    
-    # 检索相关文档
-    docs = retriever.retrieve(vector)
-    
-    # 生成答案
-    tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
-    model = RagSequenceForGeneration.from_pretrained(
-        "facebook/rag-sequence-nq",
-        torch_dtype=torch.float16
-    )
-    inputs = tokenizer.prepare_seq2seq_batch(
-        [vector], 
-        return_tensors="pt"
-    )
-    outputs = model.generate(input_ids=inputs["input_ids"])
-    return tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+    """处理向量输入并返回答案"""
+    try:
+        # start_time = time.time()
+        # load_resources()  # 确保资源已加载
+        
+        # # 转换为numpy数组
+        # vector = np.array(vector, dtype=np.float32).reshape(1, -1)
+        
+        # # 检索相关文档
+        # docs = retriever.retrieve(vector)
+        
+        # # 提取前3个相关文档
+        # context = "\n".join([doc["text"] for doc in docs[:3]])
+        
+        # # 生成答案 (使用更轻量级的生成模型)
+        # inputs = tokenizer(
+        #     f"基于以下信息回答问题: {context}\n问题: 用户查询向量",
+        #     return_tensors="pt"
+        # )
+        
+        # # 使用轻量级生成模型
+        # from transformers import AutoModelForCausalLM
+        # generator = AutoModelForCausalLM.from_pretrained("gpt2")
+        # outputs = generator.generate(
+        #     inputs["input_ids"],
+        #     max_length=200,
+        #     num_return_sequences=1
+        # )
+        
+        # answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # print(f"处理时间: {time.time() - start_time:.2f}秒")
+        # return answer
+
+        # 如果遇到资源瓶颈，使用纯检索方案
+        vector = np.array(vector, dtype=np.float32).reshape(1, -1)
+        
+        # FAISS 搜索
+        D, I = index.search(vector, k=3)
+        
+        # 获取最相关结果
+        result = metadata.iloc[I[0][0]]
+        return f"最相关结果: {result['title']}\n描述: {result['description'][:100]}..."
+    except Exception as e:
+        return f"处理错误: {str(e)}"
 
-# 创建更简单的接口
+# 创建简化接口
 with gr.Blocks() as demo:
-    gr.Markdown("## 🛍️ 电商智能客服系统")
-    
-    # 添加模型状态显示
-    model_status = gr.Markdown(f"模型状态: {'已加载' if model else '未加载'}")
+    gr.Markdown("## 🛍️ 电商智能客服系统 (轻量版)")
     
     with gr.Row():
         vector_input = gr.Dataframe(
-            headers=["vector"], 
+            headers=["向量值"], 
             type="array",
-            label="输入向量 (384维)"
+            label="输入向量 (384维)",
+            value=[[0.1]*384]  # 默认值
         )
         output = gr.Textbox(label="智能回答")
     
@@ -63,26 +130,14 @@ with gr.Blocks() as demo:
         outputs=output
     )
 
-# 在应用启动时下载索引
-INDEX_PATH = hf_hub_download(
-    repo_id="GOGO198/GOGO_rag_index",
-    filename="faiss_index.bin",
-    cache_dir="/data"
-)
-
-METADATA_PATH = hf_hub_download(
-    repo_id="GOGO198/GOGO_rag_index",
-    filename="metadata.csv",
-    cache_dir="/data"
-)
-
-# 加载索引
-index = faiss.read_index(INDEX_PATH)
-metadata = pd.read_csv(METADATA_PATH)
-    
 # 启动应用
-demo.launch(
-    server_name="0.0.0.0",
-    server_port=7860,
-    share=False
-)
+if __name__ == "__main__":
+    # 先加载必要资源
+    print("启动前预加载资源...")
+    load_resources()
+    
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )