Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import numpy as np
 import os
 from sentence_transformers import SentenceTransformer
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
 
 # 使用更轻量的模型 - 添加 token 参数
 model_name = "all-MiniLM-L6-v2"
@@ -9,24 +10,25 @@ token = os.getenv("HF_TOKEN")  # 从环境变量获取令牌
 model = SentenceTransformer(model_name, use_auth_token=token) if token else None
 
 def predict(vector):
-    """处理向量输入并返回答案"""
-    try:
-        # 检查模型是否加载成功
-        if model is None:
-            return "错误：模型未加载，请检查 HF_TOKEN 设置"
-            
-        # 转换为 numpy 数组
-        vector = np.array(vector).reshape(1, -1)
-        
-        # 验证向量维度
-        if vector.shape[1] == 384:  # MiniLM 向量维度
-            # 实际向量处理逻辑
-            return "这是基于向量生成的回答示例。实际使用时应连接您的 RAG 模型"
-        else:
-            return f"错误：接收的向量维度是 {vector.shape[1]}，但预期是 384"
-            
-    except Exception as e:
-        return f"处理错误: {str(e)}"
+    # 加载本地索引
+    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")
+    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]
 
 # 创建更简单的接口
 with gr.Blocks() as demo: