faiss向量数据库实现rag

1）环境准备

!pip install sentence-transformers transformers faiss-gpu

2）代码实现

import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# 1. 加载 SentenceTransformer 模型，用于生成嵌入
model = SentenceTransformer('paraphrase-MiniLM-L6-v2')

# 示例文档
documents = [
    "Milvus is an open-source vector database",
    "RAG is a technique to enhance generation using retrieval",
    "Python is a popular programming language",
    "Machine learning enables models to learn from data",
    "Faiss is a library for efficient similarity search and clustering of dense vectors"
]

# 2. 对文档生成嵌入向量
embeddings = model.encode(documents)
embeddings = np.array(embeddings).astype('float32')

# 3. 使用 Faiss 创建索引（L2 距离）
dimension = embeddings.shape[1]  # 嵌入向量的维度
index = faiss.IndexFlatL2(dimension)
index.add(embeddings)

# 4. 输入查询并进行向量化
query = "What is Faiss used for?"
query_embedding = model.encode([query]).astype('float32')

# 5. 在 Faiss 中检索最相似的文档
k = 3  # 查找前3个最相似的文档
distances, indices = index.search(query_embedding, k)

# 6. 获取最相似的文档
print("Top results:")
for i, idx in enumerate(indices[0]):
    print(f"Document {idx}: {documents[idx]}, Score: {distances[0][i]}")

# 7. 加载 GPT-2 模型，用于生成增强
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
gpt2_model = GPT2LMHeadModel.from_pretrained("gpt2")

# 将检索到的文档作为上下文，传递给 GPT-2 模型进行生成
context = " ".join([documents[idx] for idx in indices[0]])
input_text = f"Context: {context}\nQuestion: {query}"

inputs = tokenizer.encode(input_text, return_tensors="pt")
output = gpt2_model.generate(inputs, max_length=100, num_return_sequences=1)

generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
print("\nGenerated Response:", generated_text)

3）读取pdf和docx

!pip install PyMuPDF python-docx

import fitz  # PyMuPDF
import docx

# 读取 PDF 文件
def read_pdf(file_path):
    pdf_text = ""
    with fitz.open(file_path) as doc:
        for page_num in range(len(doc)):
            page = doc.load_page(page_num)
            pdf_text += page.get_text("text")
    return pdf_text

# 读取 DOCX 文件
def read_docx(file_path):
    doc = docx.Document(file_path)
    doc_text = "\n".join([para.text for para in doc.paragraphs])
    return doc_text

# 使用你提供的文件路径
pdf_file = "/content/cgft-llm/llama-index/pdf/Subclass 500 Student visa.pdf"  # 替换为你的 PDF 文件路径
docx_file = "/content/cgft-llm/llama-index/docs/3.98万L4无人车来了！卷出行业新低，1小时卖掉半年产能.txt"  # 替换为你的 DOCX 文件路径

# 读取文件内容
pdf_text = read_pdf(pdf_file)
# docx_text = read_docx(docx_file)

# 将读取的内容作为 documents 列表
# documents = [pdf_text, docx_text]

# 打印读取的内容（可以选择打印部分内容）
print("PDF 内容:", pdf_text[:200])  # 打印前200个字符
# print("DOCX 内容:", docx_text[:200])  # 打印前200个字符

读取txt

# 读取 TXT 文件
def read_txt(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        return file.read()

# 修改路径
txt_file = "/content/cgft-llm/llama-index/docs/3.98万L4无人车来了！卷出行业新低，1小时卖掉半年产能.txt"

# 读取文件内容
txt_text = read_txt(txt_file)



# 打印读取的内容（可以选择打印部分内容）

print("TXT 内容:", txt_text[:200])  # 打印前200个字符