34：Rerank与Query Rewrite的自动化Pipeline

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# 加载小模型
tokenizer = AutoTokenizer.from_pretrained("t5-small")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-small")

# Query Rewrite
def query_rewrite(original_query):
    """优化用户查询"""
    prompt = f"Rewrite the following query to improve search results: {original_query}"
    
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=128)
    
    rewritten_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return rewritten_query

# 示例查询
original_query = "如何使用RAG"
rewritten_query = query_rewrite(original_query)
print(f"原始查询: {original_query}")
print(f"优化后查询: {rewritten_query}")

from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.document_loaders import PyPDFLoader

# 加载文档
loader = PyPDFLoader("document.pdf")
documents = loader.load()

# 初始化嵌入模型
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

# 创建向量存储
vectorstore = Chroma.from_documents(
    documents=documents,
    embedding=embeddings
)

# 初步检索
def initial_search(query, k=10):
    """初步检索"""
    results = vectorstore.similarity_search(
        query=query,
        k=k
    )
    return results

# 示例查询
query = "如何使用RAG"
results = initial_search(query)
print("初步检索结果:")
for i, result in enumerate(results):
    print(f"{i+1}. {result.page_content[:100]}...")

from sentence_transformers import CrossEncoder

# 加载Cross-Encoder模型
model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

# Rerank
def rerank(query, documents, k=5):
    """重新排序"""
    # 准备输入对
    pairs = [(query, doc.page_content) for doc in documents]
    
    # 计算相关性得分
    scores = model.predict(pairs)
    
    # 按得分排序
    sorted_docs = sorted(
        zip(documents, scores),
        key=lambda x: x[1],
        reverse=True
    )
    
    # 返回前k个结果
    return [doc for doc, _ in sorted_docs[:k]]

# 示例查询
query = "如何使用RAG"
documents = initial_search(query)
reranked_docs = rerank(query, documents)
print("重新排序结果:")
for i, result in enumerate(reranked_docs):
    print(f"{i+1}. {result.page_content[:100]}...")

from langchain.agents import AgentType, initialize_agent, load_tools
from langchain.chat_models import ChatOpenAI

# 初始化LLM
llm = ChatOpenAI(model="gpt-4", temperature=0)

# 加载工具
tools = load_tools(["serpapi", "llm-math"], llm=llm)

# 初始化Agent
def create_agent():
    """创建Agent"""
    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True
    )
    return agent

# Agentic判断
def agentic_judgment(query, documents, agent):
    """Agentic判断"""
    # 构建上下文
    context = "\n".join([doc.page_content[:200] for doc in documents])
    
    # 构建提示
    prompt = f"基于以下上下文，判断哪些文档最相关于查询 '{query}'，并给出理由：\n\n{context}"
    
    # 执行Agent
    result = agent.run(prompt)
    
    return result

# 示例查询
agent = create_agent()
query = "如何使用RAG"
documents = rerank(query, initial_search(query))
judgment = agentic_judgment(query, documents, agent)
print("Agentic判断结果:")
print(judgment)

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.document_loaders import PyPDFLoader
from sentence_transformers import CrossEncoder
from langchain.agents import AgentType, initialize_agent, load_tools
from langchain.chat_models import ChatOpenAI

# 加载文档
loader = PyPDFLoader("document.pdf")
documents = loader.load()

# 初始化嵌入模型
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

# 创建向量存储
vectorstore = Chroma.from_documents(
    documents=documents,
    embedding=embeddings
)

# 加载Query Rewrite模型
tokenizer = AutoTokenizer.from_pretrained("t5-small")
rewrite_model = AutoModelForSeq2SeqLM.from_pretrained("t5-small")

# 加载Cross-Encoder模型
rerank_model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

# 初始化LLM和Agent
llm = ChatOpenAI(model="gpt-4", temperature=0)
tools = load_tools(["serpapi", "llm-math"], llm=llm)
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)

# Query Rewrite
def query_rewrite(original_query):
    """优化用户查询"""
    prompt = f"Rewrite the following query to improve search results: {original_query}"
    
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = rewrite_model.generate(**inputs, max_length=128)
    
    rewritten_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return rewritten_query

# 初步检索
def initial_search(query, k=10):
    """初步检索"""
    results = vectorstore.similarity_search(
        query=query,
        k=k
    )
    return results

# Rerank
def rerank(query, documents, k=5):
    """重新排序"""
    # 准备输入对
    pairs = [(query, doc.page_content) for doc in documents]
    
    # 计算相关性得分
    scores = rerank_model.predict(pairs)
    
    # 按得分排序
    sorted_docs = sorted(
        zip(documents, scores),
        key=lambda x: x[1],
        reverse=True
    )
    
    # 返回前k个结果
    return [doc for doc, _ in sorted_docs[:k]]

# Agentic判断
def agentic_judgment(query, documents, agent):
    """Agentic判断"""
    # 构建上下文
    context = "\n".join([doc.page_content[:200] for doc in documents])
    
    # 构建提示
    prompt = f"基于以下上下文，判断哪些文档最相关于查询 '{query}'，并给出理由：\n\n{context}"
    
    # 执行Agent
    result = agent.run(prompt)
    
    return result

# 完整Pipeline
def rag_pipeline(original_query):
    """完整的RAG Pipeline"""
    # 1. Query Rewrite
    rewritten_query = query_rewrite(original_query)
    print(f"原始查询: {original_query}")
    print(f"优化后查询: {rewritten_query}")
    
    # 2. 初步检索
    initial_results = initial_search(rewritten_query)
    print("初步检索结果:")
    for i, result in enumerate(initial_results[:3]):
        print(f"{i+1}. {result.page_content[:100]}...")
    
    # 3. Rerank
    reranked_results = rerank(rewritten_query, initial_results)
    print("重新排序结果:")
    for i, result in enumerate(reranked_results):
        print(f"{i+1}. {result.page_content[:100]}...")
    
    # 4. Agentic判断
    judgment = agentic_judgment(rewritten_query, reranked_results, agent)
    print("Agentic判断结果:")
    print(judgment)
    
    return reranked_results, judgment

# 示例查询
original_query = "如何使用RAG"
results, judgment = rag_pipeline(original_query)

# app.py
from fastapi import FastAPI, Query
from rag_pipeline import RAGPipeline

app = FastAPI()

# 初始化RAG Pipeline
pipeline = RAGPipeline(
    vector_store_path="./vector_store",
    rewrite_model="t5-small",
    rerank_model="cross-encoder/ms-marco-MiniLM-L-6-v2",
    llm_model="gpt-4"
)

@app.get("/search")
async def search(
    query: str = Query(..., description="搜索查询"),
    k: int = Query(5, description="返回结果数量"),
    use_rewrite: bool = Query(True, description="是否使用Query Rewrite"),
    use_rerank: bool = Query(True, description="是否使用Rerank"),
    use_agent: bool = Query(True, description="是否使用Agentic判断")
):
    """RAG Pipeline API"""
    results, judgment = pipeline.run(
        query=query,
        k=k,
        use_rewrite=use_rewrite,
        use_rerank=use_rerank,
        use_agent=use_agent
    )
    
    return {
        "query": query,
        "results": [
            {
                "content": result.page_content[:500],
                "metadata": result.metadata
            }
            for result in results
        ],
        "judgment": judgment
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

from rag_pipeline import RAGPipeline

# 初始化RAG Pipeline
pipeline = RAGPipeline(
    vector_store_path="./academic_vector_store",
    rewrite_model="t5-small",
    rerank_model="cross-encoder/ms-marco-MiniLM-L-6-v2",
    llm_model="gpt-4"
)

# 学术论文检索
query = "transformer模型在NLP中的应用"
results, judgment = pipeline.run(query=query, k=5)

print("学术论文检索结果:")
for i, result in enumerate(results):
    print(f"{i+1}. {result.page_content[:200]}...")

print("\nAgentic判断:")
print(judgment)

from rag_pipeline import RAGPipeline

# 初始化RAG Pipeline
pipeline = RAGPipeline(
    vector_store_path="./tech_vector_store",
    rewrite_model="t5-small",
    rerank_model="cross-encoder/ms-marco-MiniLM-L-6-v2",
    llm_model="gpt-4"
)

# 技术文档检索
query = "Python中的装饰器实现"
results, judgment = pipeline.run(query=query, k=5)

print("技术文档检索结果:")
for i, result in enumerate(results):
    print(f"{i+1}. {result.page_content[:200]}...")

print("\nAgentic判断:")
print(judgment)

from rag_pipeline import RAGPipeline

# 初始化RAG Pipeline
pipeline = RAGPipeline(
    vector_store_path="./qa_vector_store",
    rewrite_model="t5-small",
    rerank_model="cross-encoder/ms-marco-MiniLM-L-6-v2",
    llm_model="gpt-4"
)

# 问答数据检索
query = "如何解决Python中的内存泄漏问题"
results, judgment = pipeline.run(query=query, k=5)

print("问答数据检索结果:")
for i, result in enumerate(results):
    print(f"{i+1}. {result.page_content[:200]}...")

print("\nAgentic判断:")
print(judgment)

# 安装依赖
pip install langchain openai chromadb sentence-transformers transformers fastapi uvicorn

# 配置环境变量
export OPENAI_API_KEY=your-api-key

# 运行示例
python app.py