milvus+flask山寨《从零构建向量数据库》第7章case2

继续流水账完这本书，这个案例是打造文字形式的个人知识库雏形。

create_context_db:

# Milvus Setup Arguments
COLLECTION_NAME = 'text_content_search'
DIMENSION = 2048
MILVUS_HOST = "localhost"
MILVUS_PORT = "19530"# Inference Arguments
BATCH_SIZE = 128from pymilvus import MilvusClient,utility,connections
milvus_client = MilvusClient(uri="http://localhost:19530")# Connect to the instance
connections.connect(host=MILVUS_HOST,port=MILVUS_PORT)from markdown_processor import vectorize_segments,split_html_into_segmentstest_embedding = vectorize_segments(split_html_into_segments("<h1>RAG还是挺有意思的！</h1>"))
embedding_dim = len(test_embedding[0]) #原始的test_embedding的len结构是[[],[]]的形式
print(embedding_dim)
print(test_embedding[:10])# Remove any previous collection with the same name
if utility.has_collection(COLLECTION_NAME):utility.drop_collection(COLLECTION_NAME)milvus_client.create_collection(collection_name=COLLECTION_NAME,dimension=embedding_dim,metric_type="IP",  # Inner product distanceconsistency_level="Strong",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/consistency.md#Consistency-Level for more details.
)#下面这个手法可以直接读取md文件，然后向量化存系统。
#from tqdm import tqdm
#data = []
#from glob import glob
#text_lines = []
#for file_path in glob("milvus_docs/en/faq/*.md", recursive=True):
#    with open(file_path, "r") as file:
#        file_text = file.read()
#    text_lines += file_text.split("# ")
#
#for i, line in enumerate(tqdm(text_lines, desc="Creating embeddings")):
#    data.append({"id": i, "vector": vectorize_segments(split_html_into_segments(line)), "text": line})
#
#milvus_client.insert(collection_name=COLLECTION_NAME, data=data)

markdown_processor.py 这个文件如今大可不必了。

import markdown
from bs4 import BeautifulSoup  #用于解析和操作HTML文档
from transformers import AutoTokenizer,AutoModel #用于自动加载预训练的模型以及分词器
import torch #用于深度学习计算def markdown_to_html(markdown_text):return markdown.markdown(markdown_text)def split_html_into_segments(html_text): #定义函数，将HTML文档分割成多个段落soup = BeautifulSoup(html_text,"html.parser") #解析HTML文档segments = [] #初始化一个列表用于存储分割后的段落#找HTML文档中的段落，标题，无序列表和有序列表标签for tag in soup.find_all(["h1","h2","h3","h4","h5","h6","p","ul","ol"]):segments.append(tag.get_text())return segments#定义函数，用于将文本段落转换为向量表示
def vectorize_segments(segments):# 使用预训练的分词器和模型，这里使用的是BAAI/bge-large-zh-v1.5 一个中文模型tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-large-zh-v1.5")model = AutoModel.from_pretrained("BAAI/bge-large-zh-v1.5")model.eval() #将模型定位评估模式，避免dropout等训练模式下的参数#使用分词器对文本段落进行编码，添加必要的填充和截断，并返回PyTorch张量格式encoded_input = tokenizer(segments,padding=True,truncation=True,return_tensors="pt")with torch.no_grad():model_output = model(**encoded_input) #将编码后的输入传递给模型，获取模型的输出sentence_embeddings = model_output[0][:,0] #从模型输出中提取句子向量化的结果#对句子的量化结果进行L2归一化，以便于后续的相似度比较或聚类分析sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings,p=2,dim=1)return sentence_embeddings

flask的app.py

from flask import Flask,request,jsonify
from flask import render_template
import requests
from markdown_processor import markdown_to_html, split_html_into_segments, vectorize_segments
from pymilvus import MilvusClientimport logging
import osMILVUS_HOST = "localhost"
MILVUS_PORT = "19530"
COLLECTION_NAME = 'text_content_search'
TOP_K = 3app = Flask(__name__)
milvus_client = MilvusClient(uri="http://localhost:19530")@app.route("/")
def index():return render_template("index.html")@app.route('/upload', methods=['POST'])
def upload():if 'file' not in request.files:return jsonify({'error': 'No file part in the request'}), 400file = request.files['file']if file.filename == '':return jsonify({'error': 'No file selected for uploading'}), 400markdown_text = file.read().decode('utf-8')html_text = markdown_to_html(markdown_text)segments = split_html_into_segments(html_text)vectors = vectorize_segments(segments)# 将向量上传到数据库data = []for i, (segment, vector) in enumerate(zip(segments, vectors)):data.append({"id": i + 1,"vector": vector.tolist(), "text": segment})milvus_client.insert(collection_name=COLLECTION_NAME, data=data)return jsonify({'message': '文件已处理并上传向量到数据库'})@app.route('/search', methods=['POST'])
def search():data = request.get_json()search_text = data.get('search')# 添加前缀到查询字符串instruction = "为这个句子生成表示以用于检索相关文章："search_text_with_instruction = instruction + search_text# 向量化修改后的查询search_vector = vectorize_segments([search_text_with_instruction])[0].tolist()search_results = milvus_client.search(collection_name=COLLECTION_NAME,data=[search_vector],  limit=3,  # Return top 3 resultssearch_params={"metric_type": "IP", "params": {}},  # Inner product distanceoutput_fields=["text"],  # Return the text field)  # 构建与 LLM API 交互的消息列表messages = [{"role": "system", "content": "You are a helpful assistant. Answer questions based solely on the provided content without making assumptions or adding extra information."}] # 解析搜索结果for index,value in enumerate(search_results):#print(value)text = value[0]["entity"]["text"]print(text)messages.append({"role": "assistant", "content": text})messages.append({"role": "user", "content": search_text})# 向 deepseek 发送请求并获取答案 (用的silicon flow)url = "https://api.ap.siliconflow.com/v1/chat/completions"payload = {"model": "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B","messages": messages,"stream": False,"max_tokens": 1000,"stop": None,"temperature": 0.7,"top_p": 0.7,"top_k": 10,"frequency_penalty": 0.5,"n": 1,"response_format": {"type": "text"},}headers = {"Authorization": "Bearer <#你自己的token>","Content-Type": "application/json"}response = requests.request("POST", url, json=payload, headers=headers)answer = response.textreturn jsonify({'answer': answer})if __name__ == '__main__':app.run(host='0.0.0.0', port=5020, debug=True)

吐槽一下，silicon flow这种deepseek API免费问不到几个，就开始算钱咯。