week10 - 技术栈

目标

基于手头文件写了一个问答系统，不需要网页，能用代码做出回答就行，方法任选即可

解读

读取本地文件 → 文本切分 → 向量化 → 检索相关内容 → 调用大模型生成答案。使用RAG方法。

流程：

python 复制代码

本地文件
   ↓
Document Parsing
   ↓
文本切分 Chunk
   ↓
Embedding
   ↓
向量数据库 / FAISS
   ↓
用户问题
   ↓
Query Embedding
   ↓
Top-K Retrieval
   ↓
相关上下文
   ↓
LLM
   ↓
最终答案

方案

python 复制代码

Python
+
Sentence Transformers
+
FAISS
+
vLLM

内容

python 复制代码

pip install sentence-transformers faiss-cpu openai pypdf

import pickle

import faiss
import numpy as np

from openai import OpenAI

from sentence_transformers import (
    SentenceTransformer
)


EMBEDDING_MODEL = (
    "BAAI/bge-small-zh-v1.5"
)


LLM_MODEL = (
    "Qwen/Qwen2.5-1.5B-Instruct"
)


TOP_K = 3


embedding_model = (
    SentenceTransformer(
        EMBEDDING_MODEL
    )
)


index = faiss.read_index(
    "index/faiss.index"
)


with open(
    "index/chunks.pkl",
    "rb"
) as f:

    chunks = pickle.load(
        f
    )


client = OpenAI(
    base_url=(
        "http://localhost:8000/v1"
    ),
    api_key="EMPTY"
)


def retrieve(
    question,
    top_k=3
):

    query_embedding = (
        embedding_model.encode(
            [question],
            normalize_embeddings=True
        )
    )


    query_embedding = np.asarray(
        query_embedding,
        dtype="float32"
    )


    scores, indices = (
        index.search(
            query_embedding,
            top_k
        )
    )


    results = []


    for score, idx in zip(
        scores[0],
        indices[0]
    ):

        results.append({
            "score": float(score),
            "chunk_id": int(idx),
            "text": chunks[idx]
        })


    return results


def answer_question(
    question
):

    results = retrieve(
        question,
        TOP_K
    )


    context = "\n\n".join(
        [
            (
                f"[资料 {i + 1}] "
                f"Chunk {item['chunk_id']}\n"
                f"{item['text']}"
            )

            for i, item
            in enumerate(results)
        ]
    )


    prompt = f"""
你是一个本地文件问答助手。

请严格根据下面提供的资料回答问题。

要求：

1. 只能依据提供的资料回答。
2. 不要编造资料中不存在的信息。
3. 如果资料中无法找到答案，请回答：
   根据当前文件内容无法确定。
4. 回答尽量简洁、清晰。
5. 如果可以，请说明答案对应的资料编号。

资料：

{context}

问题：

{question}
"""


    response = (
        client.chat.completions.create(
            model=LLM_MODEL,
            messages=[
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            temperature=0,
            max_tokens=512
        )
    )


    answer = (
        response
        .choices[0]
        .message
        .content
    )


    return (
        answer,
        results
    )


if __name__ == "__main__":

    print(
        "=" * 50
    )

    print(
        "RAG 本地文件问答系统"
    )

    print(
        "输入 exit 退出"
    )

    print(
        "=" * 50
    )


    while True:

        question = input(
            "\n请输入问题："
        ).strip()


        if question.lower() in [
            "exit",
            "quit"
        ]:

            print(
                "程序退出"
            )

            break


        answer, results = (
            answer_question(
                question
            )
        )


        print(
            "\n答案："
        )

        print(
            answer
        )


        print(
            "\n检索结果："
        )


        for i, item in enumerate(
            results
        ):

            print(
                f"\nTop {i + 1}"
            )


            print(
                f"Chunk ID: "
                f"{item['chunk_id']}"
            )


            print(
                f"Similarity: "
                f"{item['score']:.4f}"
            )


            print(
                item["text"][:300]
            )