pip install langchain==0.2.0
pip install transformers==4.40.0
pip install sentence-transformers==3.0.0
pip install faiss-cpu==1.8.0  # 或者faiss-gpu如果有NVIDIA GPU
pip install chromadb==0.5.0

from langchain_community.llms import HuggingFacePipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

model_name = "meta-llama/Meta-Llama-3-8B-Instruct"

# 加载tokenizer和模型
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",  # 自动选择GPU或CPU
    torch_dtype="auto"
)

# 创建文本生成pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=512,
    temperature=0.7,
    top_p=0.95,
    repetition_penalty=1.15
)

# 创建LangChain可用的LLM对象
llm = HuggingFacePipeline(pipeline=pipe)

from langchain_community.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS

# 1. 加载文档（假设我们有一个knowledge_base目录存放文档）
loader = DirectoryLoader('./knowledge_base', glob="**/*.pdf")
documents = loader.load()

# 2. 分割文档为适合处理的块
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
    length_function=len,
)
docs = text_splitter.split_documents(documents)

# 3. 创建嵌入模型（使用开源的sentence-transformers）
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-mpnet-base-v2"
)

# 4. 创建向量存储并保存到本地
db = FAISS.from_documents(docs, embeddings)
db.save_local("faiss_index")

from langchain.chains import RetrievalQA
from langchain.memory import ConversationBufferMemory

# 加载之前保存的向量数据库
db = FAISS.load_local("faiss_index", embeddings)

# 创建检索器（设置返回前3个最相关文档）
retriever = db.as_retriever(search_kwargs={"k": 3})

# 创建带记忆的问答链
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    memory=memory,
    verbose=True,   # 显示详细处理过程（调试用）
)

# 测试问答功能
question = "如何配置产品的安全设置？"
result = qa_chain({"query": question})
print(result["result"])

from langchain.prompts import PromptTemplate

# 自定义提示模板提高回答质量
prompt_template = """基于以下上下文信息回答问题。如果你不知道答案，就说不知道。
上下文:
{context}

问题: {question}
请提供详细、专业的回答:"""

PROMPT = PromptTemplate(
    template=prompt_template, 
    input_variables=["context", "question"]
)

qa_chain.combine_documents_chain.llm_chain.prompt = PROMPT

# Reranker提升检索质量（需要额外安装）
pip install sentence-transformers[rerank]
from sentence_transformers import CrossEncoder

reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

def rerank_docs(query, docs):
    pairs = [(query, doc.page_content) for doc in docs]
    scores = reranker.predict(pairs)

    # Combine docs with scores and sort by score descendingly.
    scored_docs = list(zip(scores, docs))

    # Sort by score descending.
    scored_docs.sort(key=lambda x: x[0], reverse=True)

    return [doc for _, doc in scored_docs]

qa_chain.combine_documents_chain.document_prompt.template.relevance_score_fn = rerank_docs 

import gradio as gr

def respond(message, history):
    result = qa_chain({"query": message})
    return result["result"]

demo_chatbot = gr.ChatInterface(
    respond,
    title="企业知识库智能助手",
)

demo_chatbot.launch(server_name="0.0.0.0", server_port=7860)