{"repo": "aichijuzidetuzi/rag", "files": {"glm-RAG5.py": {"310c387d58b9c871f8a6461623fe2e8949b10645": {"t": 1785746403.053889, "n": 18292, "h": "7cb9aaa8d2eb5dcb", "c": "import os\nimport time\nimport json\nimport threading\nimport numpy as np\nimport torch\nimport faiss\nfrom flask import Flask, request, jsonify, render_template_string\nfrom transformers import AutoTokenizer, AutoModel\nimport requests\n\napp = Flask(__name__)\n\n\n# ------------------------\n# 升级配置参数\n# ------------------------\nclass EnhancedConfig:\n def __init__(self):\n self.data_files = [r\"Q&A.txt\"]\n self.top_k = 5\n self.similarity_threshold = 0.4\n self.batch_size = 512\n # 嵌入模型路径\n self.embedding_model_path = r\"C:/Users/cy/.cache/modelscope/hub/models/dienstag/chinese-bert-wwm-ext\"\n # DeepSeek API密钥\n self.deepseek_api_key = \"sk-0ae5870ea20141149406634adc7fdba1\"\n # RAG参数\n self.max_context_length = 1500\n self.llm_temperature = 0.3\n self.llm_max_length = 1024\n self.rag_prompt_template = \"\"\"你是古籍问答助手,请基于以下参考内容回答用户的问题。\n\n参考内容:\n{context}\n\n用户问题:{question}\n请用中文详细回答。如果没有明确相关内容,请说\"无相关资料\"。\"\"\"\n\n\nconfig = EnhancedConfig()\n\n\n# ------------------------\n# 模型加载(去除了accelerate依赖)\n# ------------------------\nclass EnhancedModels:\n def __init__(self):\n # 嵌入模型\n self.embed_device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n self.embed_tokenizer = AutoTokenizer.from_pretrained(config.embedding_model_path)\n self.embed_model = AutoModel.from_pretrained(config.embedding_model_path).to(self.embed_device)\n print(f\"✅ 嵌入模型加载成功 | 设备: {self.embed_device}\")\n\n\nmodels = EnhancedModels()\n\n\n# ------------------------\n# 增强版检索系统(保持不变)\n# ------------------------\nclass EnhancedRetriever:\n def __init__(self):\n self.index = None\n self.data = []\n self._build_index()\n\n def _build_index(self):\n all_items = []\n for file_path in config.data_files:\n items = self.parse_qa_file(file_path)\n all_items.extend(items)\n\n if all_items:\n texts = [item[\"content\"] for item in all_items]\n embeddings = self._batch_encode(texts)\n\n self.index = faiss.IndexFlatIP(embeddings.shape[1])\n faiss.normalize_L2(embeddings)\n self.index.add(embeddings)\n self.data = all_items\n print(f\"🛠 索引构建完成 | 数据量: {len(all_items)}\")\n else:\n raise ValueError(\"无有效数据构建索引\")\n\n def parse_qa_file(self, file_path):\n items = []\n question, answer = None, None\n with open(file_path, \"r\", encoding=\"utf-8\") as f:\n for line in f:\n line = line.strip()\n if line.startswith(\"问题:\"):\n question = line.replace(\"问题:\", \"\").strip()\n elif line.startswith(\"答案:\"):\n answer = line.replace(\"答案:\", \"\").strip()\n if question and answer:\n items.append({\n \"content\": f\"问题:{question}\\n答案:{answer}\",\n \"metadata\": {\"source\": file_path},\n \"question\": question,\n \"answer\": answer\n })\n question, answer = None, None\n return items\n\n def _batch_encode(self, texts):\n inputs = models.embed_tokenizer(\n texts,\n padding=True,\n truncation=True,\n max_length=512,\n return_tensors=\"pt\"\n ).to(models.embed_device)\n\n with torch.no_grad():\n outputs = models.embed_model(**inputs)\n return outputs.last_hidden_state[:, 0, :].cpu().numpy()\n\n def search(self, query, top_k=5):\n query_vec = self._batch_encode([query])\n faiss.normalize_L2(query_vec)\n\n distances, indices = self.index.search(query_vec, top_k * 2)\n results = []\n for score, idx in zip(distances[0], indices[0]):\n if idx >= 0 and score >= 0.5:\n results.append({\n \"content\": self.data[idx][\"content\"],\n \"question\": self.data[idx][\"question\"],\n \"answer\": self.data[idx][\"answer\"],\n \"score\": float(score),\n \"index\": int(idx)\n })\n\n return sorted(results, key=lambda x: x[\"score\"], reverse=True)[:top_k]\n\n\nretriever = EnhancedRetriever()\n\n\n# ------------------------\n# 生成系统(直接调用API生成)\n# ------------------------\nclass EnhancedGenerator:\n def generate(self, question):\n # 直接调用DeepSeek API,不提供上下文\n url = \"https://api.deepseek.com/chat/completions\"\n headers = {\n \"Content-Type\": \"application/json\",\n \"Authorization\": f\"Bearer {config.deepseek_api_key}\"\n }\n data = {\n \"model\": \"deepseek-chat\",\n \"messages\": [{\"role\": \"user\", \"content\": question}],\n \"temperature\": config.llm_temperature,\n \"max_tokens\": config.llm_max_length\n }\n\n try:\n response = requests.post(url, headers=headers, json=data)\n response.raise_for_status()\n result = response.json()\n answer = result[\"choices\"][0][\"message\"][\"content\"].strip()\n return answer\n except requests.RequestException as e:\n print(f\"API请求失败: {e}\")\n return \"系统调用API时发生错误\"\n\n\ngenerator = EnhancedGenerator()\n\n\n# ------------------------\n# 后端接口\n# ------------------------\n@app.route(\"/ask\", methods=[\"POST\"])\ndef enhanced_ask():\n start_time = time.time()\n data = request.get_json()\n question = data.get(\"question\", \"\").strip()\n\n if not question:\n return jsonify({\"answer\": \"问题内容为空\", \"sources\": []})\n\n try:\n # 获取搜索结果用于显示参考来源\n search_results = retriever.search(question, config.top_k)\n\n # 直接调用大模型生成答案\n answer = generator.generate(question)\n\n # 格式化参考来源\n sources = [{\n \"record_index\": res[\"index\"],\n \"content\": f\"问题:{res['question']}\\n答案:{res['answer']}\",\n \"metadata\": f\"相似度: {res['score']:.2f} | 来源文件: {os.path.basename(config.data_files[0])}\"\n } for res in search_results]\n\n return jsonify({\n \"answer\": answer,\n \"sources\": sources,\n \"time_cost\": round(time.time() - start_time, 2)\n })\n\n except Exception as e:\n print(f\"处理失败: {str(e)}\")\n return jsonify({\"answer\": \"系统处理时发生错误\", \"sources\": []})\n\n\n@app.route(\"/get_status\")\ndef get_status():\n return jsonify({\n \"doc_count\": len(retriever.data),\n \"qa_count\": len(retriever.data),\n \"sources\": list(set([os.path.basename(f) for f in config.data_files])),\n \"gpu_mem\": f\"{torch.cuda.memory_allocated() / 1e9:.1f}\" if torch.cuda.is_available() else \"N/A\",\n \"speed\": f\"{torch.cuda.memory_allocated() / 1e6:.1f}MB\" # 示例值\n })\n\n\n# ------------------------\n# 整合后的前端界面\n# ------------------------\n@app.route(\"/\")\ndef index():\n return render_template_string(r'''\n \n \n \n \n 古籍智能问答系统\n \n \n \n \n \n
\n
\n
\n \n
\n
\n 模型切换中,请稍候...\n
\n
\n
\n\n \n
\n
\n

系统配置

\n
\n \n \n
\n
\n\n
\n

知识库状态

\n
\n 加载中...\n
\n
\n\n
\n

性能监控

\n
\n GPU内存: 加载中...\n
\n
\n
\n\n \n
\n

\n 古籍文献智能问答系统\n

\n\n \n
\n
\n
当前模型
\n
-
\n
\n
\n
检索模式
\n
混合检索
\n
\n
\n
Embedding模型
\n
chinese-bert-wwm-ext
\n
\n
\n\n \n \n\n \n
\n \n \n \n
\n\n \n\n \n
\n
\n

智能回答

\n
\n
\n\n
\n

参考来源

\n
\n
\n
\n
\n\n \n \n \n ''')\n\n\n# ------------------------\n# 定时刷新线程(保持知识库更新)\n# ------------------------\ndef periodic_refresh(interval=300):\n while True:\n time.sleep(interval)\n try:\n retriever._build_index()\n print(\"✅ 知识库已刷新\")\n except Exception as e:\n print(f\"刷新失败: {e}\")\n\n\nrefresh_thread = threading.Thread(target=periodic_refresh, daemon=True)\nrefresh_thread.start()\n\nif __name__ == \"__main__\":\n # 验证测试\n sample_input = \"测试模型是否加载成功\"\n test_vector = retriever._batch_encode([sample_input])\n print(f\"嵌入测试输出shape: {test_vector.shape}\")\n\n test_output = generator.generate(\"测试生成\")\n print(f\"生成测试结果: {test_output}\")\n\n app.run(host=\"0.0.0.0\", port=5000, debug=True)"}}}}