{"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