基于 Milvus 和 BiomedBERT 的医学文献智能搜索系统

前言

随着医学研究的不断深入，文献数量呈爆炸式增长，如何快速从海量文献中提取关键信息成为一大挑战。最近，我基于 Milvus 向量数据库和 BiomedBERT 嵌入模型，开发了一个智能搜索系统，支持语义搜索和关键词匹配，并通过 Gradio 实现了一个交互式网页端。这篇博客将分享我的实现过程、每一步的依赖包以及安装方法，希望能为有类似需求的同学提供完整参考。

项目背景

我的目标是构建一个系统，能够：

1. 自动解析本地医学 PDF 文献，生成向量数据库。（假设我们已经把数据提取到本地了，可能你的是在数据库中）
2. 支持语义搜索（理解查询意图）和关键词搜索（精确匹配），以 OR 关系结合。
3. 提供网页界面，显示搜索结果并高亮关键词，支持分页和响应时间统计。

技术栈：

• Milvus ：高性能向量数据库。（参考上一篇从零搭建向量数据库）
• BiomedBERT：医学领域的预训练模型。
• Gradio：快速构建网页界面。
• PyPDF2 和 NLTK：解析 PDF 和文本切片。

实现过程

环境准备

在开始之前，需要安装 Python（建议 3.8+）和必要的依赖包。以下是全局依赖安装命令（假设使用 pip，建议在虚拟环境中操作）：

pip install pymilvus transformers torch numpy PyPDF2 nltk gradio

1. 数据预处理与存储（shengcheng.py）

功能：将本地 wendang 目录下的 PDF 文件解析、切片并存入 Milvus。

步骤：

• PDF 解析：提取文本并按句子切片。
• 嵌入生成：用 BiomedBERT 生成向量。
• 存储到 Milvus：保存 chunk 的 ID、向量和文本。

生成向量数据代码

import os
import PyPDF2
from transformers import AutoTokenizer, AutoModel
import torch
import numpy as np
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility, db
import nltk
from nltk.tokenize import sent_tokenize

# 下载必要的NLTK数据
nltk.download('punkt')
nltk.download('punkt_tab')

# 初始化模型和tokenizer
tokenizer = AutoTokenizer.from_pretrained("microsoft/BiomedNLP-BiomedBERT-large-uncased-abstract")
model = AutoModel.from_pretrained("microsoft/BiomedNLP-BiomedBERT-large-uncased-abstract")

# 配置
DATABASE_NAME = "my_medical_db"
COLLECTION_NAME = "new_medical_docs"
DIMENSION = 1024
CHUNK_SIZE = 512
WENDANG_DIR = "wendang"


# 连接Milvus并创建数据库
def connect_milvus():
    connections.connect(host='localhost', port='19530')
    if DATABASE_NAME not in db.list_database():
        db.create_database(DATABASE_NAME)
    db.using_database(DATABASE_NAME)


# 创建Milvus集合
def create_collection():
    if utility.has_collection(COLLECTION_NAME):
        utility.drop_collection(COLLECTION_NAME)

    fields = [
        FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=100),
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=DIMENSION),
        FieldSchema(name="chunk_text", dtype=DataType.VARCHAR, max_length=65535)  # 新增字段存储chunk文本
    ]
    schema = CollectionSchema(fields=fields, description="New medical documents collection")
    collection = Collection(COLLECTION_NAME, schema)

    index_params = {
        "metric_type": "L2",
        "index_type": "IVF_FLAT",
        "params": {"nlist": 16384}
    }
    collection.create_index("embedding", index_params)
    return collection


# 读取PDF并切片
def process_pdf(pdf_path):
    with open(pdf_path, 'rb') as file:
        pdf_reader = PyPDF2.PdfReader(file)
        num_pages = len(pdf_reader.pages)
        full_text = ""

        for page in range(num_pages):
            full_text += pdf_reader.pages[page].extract_text()

    sentences = sent_tokenize(full_text)
    chunks = []
    current_chunk = ""
    current_tokens = 0

    for sentence in sentences:
        tokens = tokenizer.encode(sentence, add_special_tokens=False)
        if current_tokens + len(tokens) <= CHUNK_SIZE:
            current_chunk += " " + sentence
            current_tokens += len(tokens)
        else:
            if current_chunk:
                chunks.append(current_chunk.strip())
            current_chunk = sentence
            current_tokens = len(tokens)

    if current_chunk:
        chunks.append(current_chunk.strip())

    return chunks


# 生成嵌入向量
def generate_embeddings(text_chunks):
    embeddings = []
    for chunk in text_chunks:
        inputs = tokenizer(chunk, return_tensors="pt", padding=True, truncation=True, max_length=CHUNK_SIZE)
        with torch.no_grad():
            outputs = model(**inputs)
            embedding = outputs.last_hidden_state[:, 0, :].squeeze().numpy()
            embeddings.append(embedding)
    return embeddings


# 存储到Milvus
def store_to_milvus(collection, embeddings, chunks, article_id):
    ids = [f"{article_id}_{i}" for i in range(len(embeddings))]
    vectors = [embedding.tolist() for embedding in embeddings]
    chunk_texts = chunks  # 存储原始chunk文本

    data = [ids, vectors, chunk_texts]
    collection.insert(data)
    collection.flush()
    return ids


# 获取wendang目录下的所有PDF文件
def get_pdf_files(directory):
    pdf_files = []
    for filename in os.listdir(directory):
        if filename.lower().endswith('.pdf'):
            pdf_files.append(os.path.join(directory, filename))
    return pdf_files


# 主流程
def main():
    connect_milvus()
    print("Milvus连接成功，当前数据库: ", DATABASE_NAME)

    collection = create_collection()
    print("Milvus集合创建成功: ", COLLECTION_NAME)

    pdf_files = get_pdf_files(WENDANG_DIR)
    if not pdf_files:
        print(f"目录 {WENDANG_DIR} 下未找到PDF文件")
        return

    for pdf_path in pdf_files:
        article_id = os.path.splitext(os.path.basename(pdf_path))[0]  # 直接使用文件名作为article_id
        print(f"\n处理文件: {pdf_path}")

        chunks = process_pdf(pdf_path)
        print(f"切片数量: {len(chunks)}")

        embeddings = generate_embeddings(chunks)
        print(f"生成嵌入数量: {len(embeddings)}")

        stored_ids = store_to_milvus(collection, embeddings, chunks, article_id)
        print(f"存储的ID (前5个): {stored_ids[:5]}...")


if __name__ == "__main__":
    main()

注意：

• 确保 Milvus 服务运行（可用 Docker 启动：docker run -d --name milvus_standalone -p 19530:19530 milvusdb/milvus:latest）。参照上面的引用链接如何使用向量数据库。
• 首次使用 BiomedBERT 会自动下载模型（约 2.5GB），需联网。默认在C:\Users\用户名\.cache\huggingface\hub

执行完上面的代码后正常的话就会在向量数据库中看到数据了（第一次记得重连一下才会看到）

2. 混合搜索与高亮（web_search.py）

功能：实现语义和关键词混合搜索，结果高亮并支持分页。

步骤：

• 语义搜索：用 BiomedBERT 生成查询向量，在 Milvus 中检索。
• 关键词搜索：用 Milvus 的 query 匹配关键词，按匹配数量动态得分。
• 网页展示：通过 Gradio 显示结果，关键词以红色背景高亮。

依赖包：最前面都已经安装

• transformers 和 torch：同上，已用于嵌入生成。
• pymilvus：Milvus 操作，已安装。
• gradio：构建网页界面

构建web端代码：

import re
import time
from transformers import AutoTokenizer, AutoModel
import torch
from pymilvus import connections, Collection, db
import gradio as gr

# 初始化模型和tokenizer
tokenizer = AutoTokenizer.from_pretrained("microsoft/BiomedNLP-BiomedBERT-large-uncased-abstract")
model = AutoModel.from_pretrained("microsoft/BiomedNLP-BiomedBERT-large-uncased-abstract")

# 配置
DATABASE_NAME = "my_medical_db"
COLLECTION_NAME = "new_medical_docs"
DIMENSION = 1024


# 连接Milvus
def connect_milvus():
    connections.connect(host='localhost', port='19530')
    db.using_database(DATABASE_NAME)


# 生成查询嵌入
def generate_query_embedding(query_text):
    inputs = tokenizer(query_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
        embedding = outputs.last_hidden_state[:, 0, :].squeeze().numpy().tolist()
    return embedding


# 高亮关键词（红色背景）
def highlight_keywords(text, keywords):
    for keyword in keywords:
        text = re.sub(
            f"(?i)({keyword})",
            r'<span style="background-color: #ffcccc; padding: 2px;">\1</span>',
            text
        )
    return text


# 语义搜索
def semantic_search(collection, query_text, top_k):
    query_embedding = generate_query_embedding(query_text)
    collection.load()
    search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
    results = collection.search(
        data=[query_embedding],
        anns_field="embedding",
        param=search_params,
        limit=top_k,
        output_fields=["id", "chunk_text"]
    )
    return [(result.entity.get("id"), result.entity.get("chunk_text"), 1 / (1 + result.distance)) for result in
            results[0]]


# 关键词搜索（动态得分）
def keyword_search(collection, keywords, top_k):
    collection.load()
    expr = " || ".join(f"chunk_text like '%{kw}%'" for kw in keywords)
    try:
        results = collection.query(expr=expr, limit=top_k, output_fields=["id", "chunk_text"])
        keyword_results = []
        for res in results:
            text = res["chunk_text"]
            # 计算匹配的关键词数量
            matched_keywords = sum(1 for kw in keywords if re.search(f"(?i){kw}", text))
            # 基础得分0.5，每匹配一个关键词加0.15，上限1.0
            score = min(0.5 + matched_keywords * 0.15, 1.0)
            keyword_results.append((res["id"], text, score))
        return keyword_results
    except Exception as e:
        print(f"关键词搜索失败: {e}")
        return []


# 混合检索（OR关系）
def hybrid_search(collection, query_text, keywords, top_k=50):
    start_time = time.time()

    results_dict = {}
    if query_text:
        semantic_results = semantic_search(collection, query_text, top_k)
        for id_, text, score in semantic_results:
            results_dict[id_] = {"id": id_, "text": text, "score": score}

    if keywords:
        keyword_results = keyword_search(collection, keywords, top_k)
        for id_, text, score in keyword_results:
            if id_ in results_dict:
                results_dict[id_]["score"] = max(results_dict[id_]["score"], score)
            else:
                results_dict[id_] = {"id": id_, "text": text, "score": score}

    final_results = []
    for item in results_dict.values():
        highlighted_text = highlight_keywords(item["text"], keywords) if keywords else item["text"]
        final_results.append({"id": item["id"], "text": highlighted_text, "score": item["score"]})

    results = sorted(final_results, key=lambda x: x["score"], reverse=True)
    search_time = time.time() - start_time
    return results, search_time


# 分页显示
def paginate_results(results, page, per_page):
    start_idx = (page - 1) * per_page
    end_idx = start_idx + per_page
    return results[start_idx:end_idx], len(results)


# Gradio搜索函数
def search(query_text, keywords_input, page, per_page):
    if not query_text and not keywords_input:
        return "请输入查询内容或关键词", ""

    connect_milvus()
    collection = Collection(COLLECTION_NAME)

    keywords = keywords_input.split() if keywords_input else []
    all_results, search_time = hybrid_search(collection, query_text, keywords)

    if not all_results:
        return "未找到匹配结果，请检查关键词或查询内容", f"搜索耗时: {search_time:.2f}秒"

    page_results, total_results = paginate_results(all_results, page, per_page)
    total_pages = (total_results + per_page - 1) // per_page

    output = f"查询: {query_text or '无'}, 关键词: {keywords or '无'}<br>"
    output += f"总结果数: {total_results}, 当前页: {page}/{total_pages}, 每页显示: {per_page}<br>"
    output += f"搜索耗时: {search_time:.2f}秒<br><br>"

    for result in page_results:
        output += f"ID: {result['id']}<br>"
        output += f"Score: {result['score']:.4f}<br>"
        output += f"Text: {result['text']}<br>"
        output += "-" * 50 + "<br>"

    return output, f"搜索耗时: {search_time:.2f}秒"


# Gradio界面
with gr.Blocks(title="医学文献搜索") as demo:
    gr.Markdown("# 医学文献搜索系统")

    with gr.Row():
        query_input = gr.Textbox(label="查询内容（可选）", placeholder="请输入查询内容...")
        keywords_input = gr.Textbox(label="关键词（可选，用空格分隔）", placeholder="请输入关键词...")

    with gr.Row():
        page_input = gr.Slider(minimum=1, maximum=10, value=1, step=1, label="页码")
        per_page_input = gr.Dropdown(choices=[5, 10, 20], value=10, label="每页显示数量")

    search_button = gr.Button("搜索")
    output = gr.HTML(label="搜索结果")
    time_output = gr.Textbox(label="响应时间", interactive=False)

    search_button.click(
        fn=search,
        inputs=[query_input, keywords_input, page_input, per_page_input],
        outputs=[output, time_output]
    )

demo.launch()

运行完上面的代码会生成一个web本地的端口网页

试下关键词搜索：

试下语义搜索：
可以看到查询的速度很快

总结

通过这个demo，我实现了从 PDF 解析到智能搜索的完整流程。依赖包的安装虽然繁琐，推荐我前面的全部一次性安装好。我也是刚开始尝试学习这块，欢迎大家尝试复现，互相学习。

依赖包总结

步骤	依赖包	安装命令	备注
数据预处理	PyPDF2	pip install PyPDF2	解析 PDF
	nltk	pip install nltk + 下载 punkt/punkt_tab	句子分割
	transformers	pip install transformers	BiomedBERT 模型
	torch	pip install torch	深度学习框架
	pymilvus	pip install pymilvus	Milvus 客户端
	numpy	pip install numpy	向量处理
搜索与网页端	gradio	pip install gradio	网页界面
	re, time	无需安装（Python 内置）	正则和高亮、时间计算