StructBERT-Large 单句对多句批量检索开发指南

import torch
import numpy as np
from transformers import AutoTokenizer, AutoModel
from typing import List, Dict

class BatchSemanticMatcher:
    def __init__(self, model_path: str):
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModel.from_pretrained(model_path)
        self.model.eval()
        self.model.half() # 使用半精度加速推理

    def encode_sentences(self, sentences: List[str]) -> torch.Tensor:
        """批量编码句子为向量"""
        inputs = self.tokenizer(
            sentences,
            padding=True,
            truncation=True,
            max_length=128,
            return_tensors="pt"
        )
        with torch.no_grad():
            outputs = self.model(**inputs.to(self.model.device))
            # 均值池化
            attention_mask = inputs['attention_mask']
            last_hidden_state = outputs.last_hidden_state
            input_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
            sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded, 1)
            sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
            return sum_embeddings / sum_mask

    def batch_similarity(self, query: str, candidates: List[str]) -> List[float]:
        """计算查询句与候选句列表的相似度"""
        # 编码所有句子
        all_sentences = [query] + candidates
        embeddings = self.encode_sentences(all_sentences)
        # 计算相似度
        query_embedding = embeddings[0:1]
        candidate_embeddings = embeddings[1:]
        similarities = torch.nn.functional.cosine_similarity(
            query_embedding, candidate_embeddings
        )
        return similarities.cpu().numpy().tolist()

4.2 批量处理优化策略

当处理大量文本时，我们需要考虑内存和计算效率：

class EfficientBatchMatcher(BatchSemanticMatcher):
    def __init__(self, model_path: str, batch_size: int = 32):
        super().__init__(model_path)
        self.batch_size = batch_size

    def process_large_corpus(self, query: str, corpus: List[str]) -> List[Dict]:
        """处理大规模文本语料库"""
        results = []
        # 分批处理候选文本
        for i in range(0, len(corpus), self.batch_size):
            batch_candidates = corpus[i:i + self.batch_size]
            batch_similarities = self.batch_similarity(query, batch_candidates)
            for j, similarity in enumerate(batch_similarities):
                results.append({
                    'candidate': batch_candidates[j],
                    'similarity': float(similarity),
                    'rank': 0 # 后续统一排序
                })
        # 按相似度排序
        results.sort(key=lambda x: x['similarity'], reverse=True)
        for i, item in enumerate(results):
            item['rank'] = i + 1
        return results

4.3 实时检索系统实现

对于需要实时响应的应用场景，我们可以实现一个更高效的检索系统：

import faiss
import pickle

class SemanticSearchSystem:
    def __init__(self, model_path: str):
        self.matcher = BatchSemanticMatcher(model_path)
        self.index = None
        self.corpus = []

    def build_index(self, corpus: List[str]):
        """构建向量索引加速检索"""
        self.corpus = corpus
        embeddings = self.matcher.encode_sentences(corpus).cpu().numpy()
        # 使用 FAISS 进行高效相似度搜索
        dimension = embeddings.shape[1]
        self.index = faiss.IndexFlatIP(dimension) # 内积索引，等价于余弦相似度
        faiss.normalize_L2(embeddings) # 归一化向量
        self.index.add(embeddings)

    def search(self, query: str, top_k: int = 10) -> List[Dict]:
        """语义搜索"""
        if self.index is None:
            raise ValueError("请先构建索引")
        query_embedding = self.matcher.encode_sentences([query]).cpu().numpy()
        faiss.normalize_L2(query_embedding) # 归一化向量
        # 搜索最相似的 top_k 个结果
        similarities, indices = self.index.search(query_embedding, top_k)
        results = []
        for i, idx in enumerate(indices[0]):
            results.append({
                'rank': i + 1,
                'candidate': self.corpus[idx],
                'similarity': float(similarities[0][i])
            })
        return results

5. 完整应用示例

5.1 批量检索 Web 应用

下面是一个完整的 Streamlit 应用示例，支持单句对多句的批量检索：

import streamlit as st
import pandas as pd
from batch_matcher import EfficientBatchMatcher

# 初始化模型
@st.cache_resource
def load_model():
    return EfficientBatchMatcher(
        "/root/ai-models/iic/nlp_structbert_sentence-similarity_chinese-large/"
    )

def main():
    st.title("StructBERT 批量语义检索系统")
    # 初始化模型
    matcher = load_model()
    # 输入界面
    st.sidebar.header("检索配置")
    top_k = st.sidebar.slider("返回结果数量", 5, 50, 10)
    st.header("输入查询句子")
    query = st.text_area("请输入要查询的句子", height=100)
    st.header("输入候选句子集合")
    candidate_text = st.text_area(
        "请输入候选句子，每行一个句子",
        height=200,
        help="每个候选句子单独占一行"
    )
    if st.button("开始批量检索", type="primary"):
        if not query or not candidate_text:
            st.warning("请先输入查询句子和候选句子")
            return
        candidates = [line.strip() for line in candidate_text.split('\n') if line.strip()]
        with st.spinner("正在计算相似度..."):
            results = matcher.process_large_corpus(query, candidates)
        # 显示结果
        st.header("检索结果")
        # 转换为 DataFrame 便于显示
        df_data = []
        for item in results[:top_k]:
            df_data.append({
                '排名': item['rank'],
                '候选句子': item['candidate'],
                '相似度': f"{item['similarity']:.4f}",
                '匹配程度': '高度相似' if item['similarity'] > 0.85 else '相关' if item['similarity'] > 0.5 else '不相关'
            })
        df = pd.DataFrame(df_data)
        st.dataframe(df, use_container_width=True)
        # 可视化显示
        st.subheader("相似度分布")
        similarities = [item['similarity'] for item in results[:top_k]]
        st.bar_chart(pd.DataFrame({'相似度': similarities}))

if __name__ == "__main__":
    main()

5.2 性能优化建议

对于生产环境的大规模应用，考虑以下优化策略：

# 异步处理支持
import asyncio
from concurrent.futures import ThreadPoolExecutor

class AsyncBatchMatcher:
    def __init__(self, model_path: str, max_workers: int = 4):
        self.matcher = BatchSemanticMatcher(model_path)
        self.executor = ThreadPoolExecutor(max_workers=max_workers)

    async def async_batch_similarity(self, query: str, candidates: List[str]):
        """异步批量相似度计算"""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            self.executor, self.matcher.batch_similarity, query, candidates
        )

6. 实际应用场景案例

6.1 智能客服问答匹配

在客服系统中，使用批量检索模式可以快速找到用户问题的最相关答案：

class FAQMatcher:
    def __init__(self, model_path: str):
        self.search_system = SemanticSearchSystem(model_path)
        self.faq_data = [] # 存储问答对

    def load_faqs(self, faq_list: List[Dict]):
        """加载 FAQ 数据"""
        questions = [item['question'] for item in faq_list]
        self.faq_data = faq_list
        self.search_system.build_index(questions)

    def find_best_answer(self, user_question: str, threshold: float = 0.7) -> Dict:
        """查找最相关答案"""
        results = self.search_system.search(user_question, top_k=1)
        if results and results[0]['similarity'] >= threshold:
            best_match = results[0]
            faq_index = self.search_system.corpus.index(best_match['candidate'])
            return {
                'answer': self.faq_data[faq_index]['answer'],
                'similarity': best_match['similarity'],
                'matched_question': best_match['candidate']
            }
        else:
            return {'answer': None, 'similarity': 0, 'matched_question': None}

6.2 文档内容去重检测

批量处理文档集合，检测和消除重复内容：

class DuplicateDetector:
    def __init__(self, model_path: str, similarity_threshold: float = 0.9):
        self.matcher = BatchSemanticMatcher(model_path)
        self.threshold = similarity_threshold

    def find_duplicates(self, documents: List[str]) -> List[List[int]]:
        """查找重复文档组"""
        # 编码所有文档
        embeddings = self.matcher.encode_sentences(documents)
        embeddings = embeddings.cpu().numpy()
        # 计算相似度矩阵
        from sklearn.metrics.pairwise import cosine_similarity
        sim_matrix = cosine_similarity(embeddings)
        # 查找重复组
        duplicates = []
        visited = set()
        for i in range(len(documents)):
            if i in visited:
                continue
            duplicate_group = [i]
            for j in range(i + 1, len(documents)):
                if sim_matrix[i][j] >= self.threshold:
                    duplicate_group.append(j)
                    visited.add(j)
            if len(duplicate_group) > 1:
                duplicates.append(duplicate_group)
                visited.add(i)
        return duplicates

7. 总结

通过本教程，我们深入探讨了如何基于 StructBERT-Large 模型开发单句对多句的批量检索系统。关键要点包括：

技术核心：利用 StructBERT 的深度语义理解能力，结合均值池化技术生成高质量的句子向量，通过余弦相似度实现精准的语义匹配。

扩展能力：从基础的单句对比扩展到批量处理模式，支持大规模文本检索、智能问答匹配、文档去重等多种应用场景。

性能优化：采用向量索引、批量处理、异步计算等技术手段，确保系统在大规模数据处理时仍能保持高效性能。

实践价值：这套方案可以直接应用于实际业务场景，如智能客服系统、内容检索平台、文档管理系统等，显著提升文本处理效率和准确性。

StructBERT-Large 的强大语义理解能力，加上合理的工程实现，为我们处理中文文本相似度任务提供了可靠的技术方案。通过本指南提供的代码示例和实践建议，你可以快速构建属于自己的批量语义检索系统。