Update app.py

app.py

@@ -1,15 +1,14 @@
-import gradio as gr
+ import gradio as gr
 import numpy as np
 import os
 import torch
 import pandas as pd
-from sentence_transformers import SentenceTransformer
-from huggingface_hub import hf_hub_download
 import faiss
+from huggingface_hub import hf_hub_download
+from transformers import CLIPProcessor, CLIPModel
 import time
-import pathlib
 
-# 创建安全缓存目录（在用户目录下）
+# 创建安全缓存目录
 CACHE_DIR = "/home/user/cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
 
@@ -17,105 +16,97 @@ os.makedirs(CACHE_DIR, exist_ok=True)
 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
+clip_model = None
+clip_processor = None
 
 def load_resources():
-    """按需加载资源"""
-    global model, index, metadata, tokenizer, retriever
+    """加载所有必要资源（768维专用）"""
+    global index, metadata, clip_model, clip_processor
     
-    # 仅当需要时加载
-    if model is None:
-        print("正在加载句子嵌入模型...")
-        token = os.getenv("HF_TOKEN")
-        model = SentenceTransformer("all-MiniLM-L6-v2", use_auth_token=token, cache_folder=CACHE_DIR)
-        print("句子模型加载完成")
-
+    # 加载CLIP模型（用于维度验证）
+    if clip_model is None:
+        print("正在加载CLIP模型...")
+        clip_model = CLIPModel.from_pretrained(
+            "openai/clip-vit-large-patch14",
+            cache_dir=CACHE_DIR
+        )
+        clip_processor = CLIPProcessor.from_pretrained(
+            "openai/clip-vit-large-patch14",
+            cache_dir=CACHE_DIR
+        )
+        print("CLIP模型加载完成")
+    
+    # 加载FAISS索引（768维）
     if index is None:
         print("正在下载FAISS索引...")
         INDEX_PATH = hf_hub_download(
             repo_id="GOGO198/GOGO_rag_index",
             filename="faiss_index.bin",
             cache_dir=CACHE_DIR,
-            use_auth_token=os.getenv("HF_TOKEN")
+            token=os.getenv("HF_TOKEN")
         )
         index = faiss.read_index(INDEX_PATH)
-        print("FAISS索引加载完成")
-
+        
+        # 验证索引维度
+        if index.d != 768:
+            raise ValueError(f"索引维度错误：预期768维，实际{index.d}维")
+        print("FAISS索引加载完成 | 维度: 768")
+    
+    # 加载元数据
     if metadata is None:
         print("正在下载元数据...")
         METADATA_PATH = hf_hub_download(
             repo_id="GOGO198/GOGO_rag_index",
             filename="metadata.csv",
             cache_dir=CACHE_DIR,
-            use_auth_token=os.getenv("HF_TOKEN")
+            token=os.getenv("HF_TOKEN")
         )
         metadata = pd.read_csv(METADATA_PATH)
         print("元数据加载完成")
 
 def predict(vector):
-    """处理向量输入并返回答案"""
+    """处理768维向量输入并返回答案"""
     try:
-        # start_time = time.time()
-        # load_resources()  # 确保资源已加载
-        
-        # # 转换为numpy数组
-        # vector = np.array(vector, dtype=np.float32).reshape(1, -1)
+        start_time = time.time()
+        load_resources()  # 确保资源已加载
         
-        # # 检索相关文档
-        # 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
-
-        # 如果遇到资源瓶颈，使用纯检索方案1
+        # 转换为numpy数组并验证维度
         vector = np.array(vector, dtype=np.float32).reshape(1, -1)
+        if vector.shape[1] != 768:
+            return f"维度错误：预期768维，收到{vector.shape[1]}维"
         
         # FAISS 搜索
         D, I = index.search(vector, k=3)
         
         # 获取最相关结果
-        result = metadata.iloc[I[0][0]]
-        return f"最相关结果: {result['title']}\n描述: {result['description'][:100]}..."
+        results = []
+        for i in range(3):
+            result = metadata.iloc[I[0][i]]
+            results.append(f"相关结果 {i+1}: {result['source']} | 相似度: {1/(1+D[0][i]):.2f}")
+        
+        response = "\n".join(results)
+        print(f"处理时间: {time.time() - start_time:.2f}秒")
+        return response
+        
     except Exception as e:
         return f"处理错误: {str(e)}"
 
 # 创建简化接口
 with gr.Blocks() as demo:
-    gr.Markdown("## 🛍️ 电商智能客服系统 (轻量版)")
+    gr.Markdown("## 🛍 电商智能客服系统 (768维专用)")
+    gr.Markdown("**使用CLIP-vit-large-patch14模型 | 向量维度: 768**")
     
     with gr.Row():
         vector_input = gr.Dataframe(
             headers=["向量值"], 
             type="array",
-            label="输入向量 (384维)",
-            value=[[0.1]*384]  # 默认值
+            label="输入向量 (768维)",
+            value=[[0.1]*768]  # 768维默认值
         )
-        output = gr.Textbox(label="智能回答")
+        output = gr.Textbox(label="智能回答", lines=5)
     
     submit_btn = gr.Button("生成回答")
     submit_btn.click(
@@ -127,14 +118,14 @@ with gr.Blocks() as demo:
 # 启动应用
 if __name__ == "__main__":
     # 确保缓存目录存在
+    import pathlib
     pathlib.Path(CACHE_DIR).mkdir(parents=True, exist_ok=True)
     
-    # 先加载必要资源
+    # 预加载资源
     print("启动前预加载资源...")
     load_resources()
     
     demo.launch(
         server_name="0.0.0.0",
-        server_port=7860,
-        share=False
-    )
+        server_port=7860
+    )