# 模型路径配置
model_name = "qwen/Qwen1.5-1.8B-Chat"
embedding_model_id = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
cache_dir = "D:\\modelscope\\hub"
# 加载本地大模型（CPU兼容）
local_model_path = snapshot_download(model_name, cache_dir=cache_dir)
tokenizer = AutoTokenizer.from_pretrained(local_model_path, trust_remote_code=True)
llm_model = AutoModelForCausalLM.from_pretrained(
    local_model_path, trust_remote_code=True,
    torch_dtype=torch.float32, device_map="cpu"
).eval()
# 加载向量模型
embedding_local_path = snapshot_download(embedding_model_id, cache_dir=cache_dir)
embedding_model = SentenceTransformer(embedding_local_path, device="cpu")

# --------------------------
# 3. 自定义LangChain Embedding适配器
# 作用：将句子转换为向量，用于语义匹配
# --------------------------
class LocalEmbedding(Embeddings):
    def __init__(self, model):
        self.model = model
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return self.model.encode(texts, convert_to_numpy=True).tolist()
    def embed_query(self, text: str) -> List[float]:
        return self.model.encode(text, convert_to_numpy=True).tolist()
embedding = LocalEmbedding(embedding_model)

class LocalQwenLLM(LLM):
    tokenizer: Any
    model: Any
    max_new_tokens: int = 512
    temperature: float = 0.1  # 低温度=高确定性，杜绝幻觉
    @property
    def _llm_type(self) -> str:
        return "qwen_local"
    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        inputs = self.tokenizer(prompt, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs, max_new_tokens=self.max_new_tokens,
                temperature=self.temperature, do_sample=False
            )
        return self.tokenizer.decode(outputs[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)

# 医疗专业知识库（可扩展PDF/Excel/数据库）
medical_knowledge = [
    Document(page_content="高血压饮食建议：低盐饮食，每日食盐不超过5克...", metadata={"category": "饮食", "disease": "高血压"}),
    Document(page_content="糖尿病饮食建议：控制碳水化合物摄入量...", metadata={"category": "饮食", "disease": "糖尿病"}),
    Document(page_content="感冒护理建议：保证充足睡眠，多饮温水...", metadata={"category": "护理", "disease": "感冒"}),
    Document(page_content="冠心病预防：戒烟限酒，规律运动...", metadata={"category": "预防", "disease": "冠心病"})
]
# 构建FAISS向量库 + 本地持久化
vector_store = FAISS.from_documents(medical_knowledge, embedding)
vector_store.save_local("D:\\medical_rag_faiss_db", "medical_index")

def show_similarity_scores(question: str):
    """展示问题与所有文档的相似度得分"""
    print(f"\n[向量相似度分析] 问题: {question}")
    print("-" * 50)
    # 计算问题向量
    query_vector = embedding.embed_query(question)
    # 计算与每个文档的相似度
    import numpy as np
    for i, doc in enumerate(medical_knowledge):
        doc_vector = embedding.embed_query(doc.page_content)
        # 余弦相似度
        query_np = np.array(query_vector)
        doc_np = np.array(doc_vector)
        similarity = np.dot(query_np, doc_np) / (np.linalg.norm(query_np) * np.linalg.norm(doc_np))
        # 标记是否超过阈值
        status = "✓" if similarity > 0.6 else "✗"
        print(f"{status} 文档{i+1} ({doc.metadata['disease']}): {similarity:.4f}")
        print(f"  内容: {doc.page_content[:50]}...")
    print()

def rag_query(question: str, k: int = 2, threshold: float = 0.6) -> dict:
    # 1. 计算问题向量
    query_vector = embedding.embed_query(question)
    
    # 2. 余弦相似度检索 + 阈值过滤
    scored_docs = []
    for doc in medical_knowledge:
        doc_vector = embedding.embed_query(doc.page_content)
        similarity = np.dot(query_np, doc_np) / (np.linalg.norm(query_np) * np.linalg.norm(doc_np))
        if similarity > threshold:
            scored_docs.append((doc, similarity))
    
    # 3. 按相似度排序，取Top2
    scored_docs.sort(key=lambda x: x[1], reverse=True)
    docs = [doc for doc, score in scored_docs[:k]]
    # 4. 构建约束型提示词（强制基于知识回答）
    prompt = f"""你是专业医疗顾问，严格基于知识库回答，不编造信息。
    知识库：{context}
    用户问题：{question}"""
    # 5. 大模型生成答案
    answer = local_llm._call(prompt)
    return {"question": question, "answer": answer, "source_documents": docs}

import os
import torch
import matplotlib.pyplot as plt
from modelscope import snapshot_download
from transformers import AutoTokenizer, AutoModelForCausalLM
from sentence_transformers import SentenceTransformer
# LangChain 核心模块
from langchain_core.embeddings import Embeddings
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_core.language_models.llms import LLM
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from typing import Any, List, Optional
# --------------------------
# 1. 配置路径（复用你的本地模型）
# --------------------------
model_name = "qwen/Qwen1.5-1.8B-Chat"
embedding_model_id = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
cache_dir = "D:\\modelscope\\hub"
faiss_db_path = "D:\\medical_rag_faiss_db"  # 向量库持久化路径
os.makedirs(faiss_db_path, exist_ok=True)
# --------------------------
# 2. 下载/加载本地模型（无重复下载）
# --------------------------
print("="*50)
print("加载本地大模型+Embedding模型...")
# 加载对话大模型
local_model_path = snapshot_download(model_name, cache_dir=cache_dir)
tokenizer = AutoTokenizer.from_pretrained(local_model_path, trust_remote_code=True)
llm_model = AutoModelForCausalLM.from_pretrained(
    local_model_path,
    trust_remote_code=True,
    torch_dtype=torch.float32,  # CPU兼容
    device_map="cpu"
).eval()
# 加载Embedding模型
embedding_local_path = snapshot_download(embedding_model_id, cache_dir=cache_dir, revision="master")
embedding_model = SentenceTransformer(embedding_local_path, device="cpu")
# --------------------------
# 3. 自定义LangChain Embedding适配器
# 作用：将句子转换为向量，用于语义匹配
# --------------------------
class LocalEmbedding(Embeddings):
    def __init__(self, model):
        self.model = model
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return self.model.encode(texts, convert_to_numpy=True).tolist()
    def embed_query(self, text: str) -> List[float]:
        return self.model.encode(text, convert_to_numpy=True).tolist()
embedding = LocalEmbedding(embedding_model)
# --------------------------
# 4. 自定义LangChain LLM适配器
# 作用：对接本地Qwen大模型，实现生成能力
# --------------------------
class LocalQwenLLM(LLM):
    tokenizer: Any
    model: Any
    max_new_tokens: int = 512
    temperature: float = 0.1
    @property
    def _llm_type(self) -> str:
        return "qwen_local"
    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        inputs = self.tokenizer(prompt, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=self.max_new_tokens,
                temperature=self.temperature,
                do_sample=False,
                pad_token_id=self.tokenizer.eos_token_id
            )
        response = self.tokenizer.decode(outputs[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
        return response
# 初始化本地大模型包装器
local_llm = LocalQwenLLM(tokenizer=tokenizer, model=llm_model)
# --------------------------
# 5. 构建医疗知识库（可扩展为Excel/文档读取）
# --------------------------
print("构建医疗知识库...")
medical_knowledge = [
    Document(
        page_content="高血压饮食建议：低盐饮食，每日食盐不超过5克；多吃蔬菜、全谷物；减少油炸食品、动物内脏摄入",
        metadata={"category": "饮食", "disease": "高血压"}
    ),
    Document(
        page_content="糖尿病饮食建议：控制碳水化合物摄入量，优先选择低升糖指数食物，定时定量进餐，避免含糖饮料",
        metadata={"category": "饮食", "disease": "糖尿病"}
    ),
    Document(
        page_content="感冒护理建议：保证充足睡眠，多饮温水，症状严重时可服用解热镇痛药，病程一般3-7天",
        metadata={"category": "护理", "disease": "感冒"}
    ),
    Document(
        page_content="冠心病预防：戒烟限酒，规律运动，控制血压血脂，避免情绪剧烈波动，定期体检",
        metadata={"category": "预防", "disease": "冠心病"}
    )
]
# --------------------------
# 6. 构建FAISS向量库（核心检索模块）
# 技术细节：FAISS使用L2距离进行语义相似度计算
# --------------------------
print("构建FAISS向量检索库...")
vector_store = FAISS.from_documents(medical_knowledge, embedding)
# 持久化保存向量库（下次直接加载，无需重新构建）
vector_store.save_local(faiss_db_path, "medical_index")
# --------------------------
# 7. 生成FAISS检索流程图（可视化输出图片）
# --------------------------
plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.figure(figsize=(10, 6))
plt.title("医疗知识库RAG检索流程", fontsize=16, pad=20)
# 流程节点
steps = [
    "用户问题",
    "Embedding向量化",
    "FAISS语义检索",
    "召回相关医疗知识",
    "大模型生成答案"
]
# 节点坐标
x = [1, 2, 3, 4, 5]
y = [1, 1, 1, 1, 1]
# 绘制流程
plt.plot(x, y, "o-", linewidth=3, markersize=15, color="#2E86AB")
for i, step in enumerate(steps):
    plt.text(x[i], y[i]+0.1, step, ha="center", fontsize=12, bbox=dict(boxstyle="round,pad=0.3", facecolor="#A2D6F9"))
plt.ylim(0.5, 1.5)
plt.axis("off")
plt.tight_layout()
plt.savefig("medical_rag_process.png", dpi=300, bbox_inches="tight")
print("✅ 检索流程图已保存：medical_rag_process.png")
# --------------------------
# 8. 手动构建RAG问答链（简化版，避免LangChain兼容性问题）
# 检索参数：top_k=2 召回最相关的2条知识
# --------------------------
def rag_query(question: str, k: int = 2, threshold: float = 0.6) -> dict:
    """RAG查询函数"""
    # 1. 检索相关文档并计算相似度
    import numpy as np
    # 获取所有文档的相似度
    query_vector = embedding.embed_query(question)
    scored_docs = []
    for doc in medical_knowledge:
        doc_vector = embedding.embed_query(doc.page_content)
        query_np = np.array(query_vector)
        doc_np = np.array(doc_vector)
        similarity = np.dot(query_np, doc_np) / (np.linalg.norm(query_np) * np.linalg.norm(doc_np))
        # 只保留相似度大于阈值的文档
        if similarity > threshold:
            scored_docs.append((doc, similarity))
    # 按相似度排序，取top_k
    scored_docs.sort(key=lambda x: x[1], reverse=True)
    docs = [doc for doc, score in scored_docs[:k]]
    # 打印检索到的文档（验证RAG是否生效）
    print(f"\n[调试] 检索到 {len(docs)} 条相关文档（阈值: {threshold}）：")
    for i, (doc, score) in enumerate(scored_docs[:k]):
        print(f"  文档{i+1} (相似度: {score:.4f}): {doc.page_content}")
        print(f"  元数据: {doc.metadata}")
    # 2. 判断是否有检索结果
    if len(docs) == 0:
        answer = f"抱歉，知识库中没有与您的问题相关的信息（最高相似度低于{threshold}）。"
        return {
            "question": question,
            "answer": answer,
            "source_documents": []
        }
    # 3. 构建提示词
    context = "\n\n".join([f"知识{i+1}: {doc.page_content}" for i, doc in enumerate(docs)])
    prompt = f"""你是一名专业的医疗顾问。请**严格基于以下知识库内容**回答用户问题。
【重要要求】
1. 你的答案必须来自下方提供的知识库内容
2. 不要编造知识库以外的信息
3. 如果知识库中没有相关信息，请明确说明"知识库中没有相关信息"
【知识库内容】
{context}
【用户问题】
{question}
【请基于知识库内容回答】"""
    # 4. 调用大模型生成答案
    answer = local_llm._call(prompt)
    return {
        "question": question,
        "answer": answer,
        "source_documents": docs
    }
def show_similarity_scores(question: str):
    """展示问题与所有文档的相似度得分"""
    print(f"\n[向量相似度分析] 问题: {question}")
    print("-" * 50)
    # 计算问题向量
    query_vector = embedding.embed_query(question)
    # 计算与每个文档的相似度
    import numpy as np
    for i, doc in enumerate(medical_knowledge):
        doc_vector = embedding.embed_query(doc.page_content)
        # 余弦相似度
        query_np = np.array(query_vector)
        doc_np = np.array(doc_vector)
        similarity = np.dot(query_np, doc_np) / (np.linalg.norm(query_np) * np.linalg.norm(doc_np))
        # 标记是否超过阈值
        status = "✓" if similarity > 0.6 else "✗"
        print(f"{status} 文档{i+1} ({doc.metadata['disease']}): {similarity:.4f}")
        print(f"  内容: {doc.page_content[:50]}...")
    print()
# --------------------------
# 9. 医疗问答测试（实战演示）
# --------------------------
print("\n" + "="*50)
print("医疗问答测试")
print("="*50)
# 测试问题1 - 高血压（应该能准确检索）
query1 = "高血压患者平时应该怎么吃饭？"
print(f"\n用户问题：{query1}")
result1 = rag_query(query1)
print(f"模型答案：{result1['answer']}")
# 测试问题2 - 感冒（向量检索可能不精确）
print("\n" + "="*50)
query2 = "感冒了需要注意什么？"
print(f"\n用户问题：{query2}")
print("\n[对比测试] 直接查看知识库中的感冒文档：")
for doc in medical_knowledge:
    if "感冒" in doc.page_content:
        print(f"找到：{doc.page_content}")
# 先展示向量相似度分析
show_similarity_scores("感冒了需要注意什么？")
result2 = rag_query(query2)
print(f"模型答案：{result2['answer']}")
# 测试问题3 - 直接询问感冒护理（使用更精确的关键词）
print("\n" + "="*50)
query3 = "感冒护理建议"
print(f"\n用户问题：{query3}")
show_similarity_scores("感冒护理建议")
result3 = rag_query(query3)
print(f"模型答案：{result3['answer']}")
# --------------------------
# 10. 加载已有向量库（生产环境使用）
# --------------------------
# vector_store = FAISS.load_local(faiss_db_path, embedding, allow_dangerous_deserialization=True)
# 然后重新调用rag_query函数
print("\n✅ 医疗知识库RAG系统部署完成！")