nlp_structbert_sentence-similarity_chinese-large实战教程扩展为单句vs多句批量检索的代码改造指南1. 项目简介与扩展需求nlp_structbert_sentence-similarity_chinese-large 是基于阿里达摩院开源的 StructBERT 大规模预训练模型开发的本地化语义匹配工具。这个工具能够将中文句子转化为高质量的特征向量通过余弦相似度算法精准量化句子间的语义相关性。为什么需要扩展批量检索功能在实际应用中我们经常遇到这样的场景需要将一个查询句子与成千上万个候选句子进行相似度比较然后找出最相似的几个结果。比如在知识库中搜索相关问题答案在海量文档中查找相似内容为用户查询匹配最相关的FAQ条目原来的工具只能比较两个句子现在我们要改造它让它能够处理一个句子 vs 多个句子的批量检索任务。2. 环境准备与基础代码分析2.1 环境要求确保已安装以下依赖库pip install torch transformers streamlit numpy2.2 原始代码结构分析先来看看原始代码的核心部分import torch from transformers import AutoTokenizer, AutoModel import streamlit as st # 模型加载使用缓存避免重复加载 st.cache_resource def load_model(): model_path /root/ai-models/iic/nlp_structbert_sentence-similarity_chinese-large tokenizer AutoTokenizer.from_pretrained(model_path) model AutoModel.from_pretrained(model_path).half().cuda() return tokenizer, model # 获取句子向量的函数 def get_sentence_embedding(text, tokenizer, model): inputs tokenizer(text, return_tensorspt, paddingTrue, truncationTrue, max_length512) with torch.no_grad(): outputs model(**inputs.to(cuda)) attention_mask inputs[attention_mask] token_embeddings outputs.last_hidden_state input_mask_expanded attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float().to(cuda) sum_embeddings torch.sum(token_embeddings * input_mask_expanded, 1) sum_mask torch.clamp(input_mask_expanded.sum(1), min1e-9) return sum_embeddings / sum_mask # 计算余弦相似度 def cosine_similarity(embedding1, embedding2): return torch.nn.functional.cosine_similarity(embedding1, embedding2)3. 批量检索功能改造3.1 核心改造思路原来的代码一次只处理两个句子的比较现在我们要改造为一次处理一个查询句子和多个候选句子批量计算相似度避免重复计算返回排序后的相似度结果3.2 批量处理函数实现def batch_similarity_search(query_sentence, candidate_sentences, tokenizer, model, top_k5): 批量相似度搜索函数 参数: - query_sentence: 查询句子字符串 - candidate_sentences: 候选句子列表 - tokenizer: 模型tokenizer - model: 模型 - top_k: 返回最相似的几个结果 返回: - 排序后的相似度结果列表 # 获取查询句子的向量 query_embedding get_sentence_embedding(query_sentence, tokenizer, model) # 批量获取候选句子的向量 candidate_embeddings [] for sentence in candidate_sentences: embedding get_sentence_embedding(sentence, tokenizer, model) candidate_embeddings.append(embedding) # 将候选向量堆叠成一个张量 candidate_embeddings torch.cat(candidate_embeddings, dim0) # 批量计算相似度 similarities cosine_similarity(query_embedding, candidate_embeddings) # 组合结果并排序 results [] for i, similarity in enumerate(similarities): results.append({ sentence: candidate_sentences[i], similarity: similarity.item(), rank: 0 }) # 按相似度降序排序 results.sort(keylambda x: x[similarity], reverseTrue) # 添加排名 for rank, result in enumerate(results[:top_k], 1): result[rank] rank return results[:top_k]3.3 优化版批量处理上面的代码虽然能用但效率不够高。我们来优化一下def optimized_batch_search(query_sentence, candidate_sentences, tokenizer, model, top_k5, batch_size32): 优化版的批量搜索函数支持分批次处理大量数据 参数: - batch_size: 每次处理的句子数量避免内存溢出 # 获取查询句子的向量 query_embedding get_sentence_embedding(query_sentence, tokenizer, model) all_similarities [] # 分批次处理候选句子 for i in range(0, len(candidate_sentences), batch_size): batch_sentences candidate_sentences[i:ibatch_size] # 批量编码 inputs tokenizer( batch_sentences, return_tensorspt, paddingTrue, truncationTrue, max_length512 ) with torch.no_grad(): outputs model(**inputs.to(cuda)) # 批量计算均值池化 attention_mask inputs[attention_mask] token_embeddings outputs.last_hidden_state input_mask_expanded attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float().to(cuda) sum_embeddings torch.sum(token_embeddings * input_mask_expanded, 1) sum_mask torch.clamp(input_mask_expanded.sum(1), min1e-9) batch_embeddings sum_embeddings / sum_mask # 批量计算相似度 batch_similarities cosine_similarity( query_embedding.repeat(len(batch_sentences), 1), batch_embeddings ) all_similarities.extend(batch_similarities.cpu().numpy()) # 组合结果 results [] for i, similarity in enumerate(all_similarities): results.append({ sentence: candidate_sentences[i], similarity: float(similarity), rank: 0 }) # 排序并返回top_k results.sort(keylambda x: x[similarity], reverseTrue) for rank, result in enumerate(results[:top_k], 1): result[rank] rank return results[:top_k]4. Streamlit界面改造4.1 新的用户界面现在我们来改造Streamlit界面支持批量检索功能import streamlit as st import pandas as pd # 页面设置 st.set_page_config(page_titleStructBERT 批量语义检索, layoutwide) # 加载模型 tokenizer, model load_model() st.title( StructBERT 中文句子批量语义检索) # 创建两列布局 col1, col2 st.columns([1, 2]) with col1: st.header(输入查询) # 查询句子输入 query_sentence st.text_area( 请输入查询句子, height100, placeholder例如手机的电池续航时间多久 ) # 候选句子输入 candidate_text st.text_area( 请输入候选句子每行一个, height300, placeholder例如这款手机电池很耐用\n手机续航能力强大\n电池可以使用一整天\n... ) top_k st.slider(返回最相似的前几个结果, 1, 20, 5) if st.button( 开始批量检索, typeprimary): if query_sentence and candidate_text: candidate_sentences [s.strip() for s in candidate_text.split(\n) if s.strip()] if candidate_sentences: with st.spinner(f正在处理 {len(candidate_sentences)} 个候选句子...): results optimized_batch_search( query_sentence, candidate_sentences, tokenizer, model, top_k ) st.session_state.results results st.session_state.query query_sentence else: st.warning(请输入至少一个候选句子) else: st.warning(请填写查询句子和候选句子) with col2: st.header(检索结果) if results in st.session_state: st.subheader(f查询句子{st.session_state.query}) # 创建结果表格 result_data [] for result in st.session_state.results: similarity result[similarity] color green if similarity 0.85 else orange if similarity 0.5 else red result_data.append({ 排名: result[rank], 候选句子: result[sentence], 相似度: f{similarity:.4f}, 语义关系: 非常相似 if similarity 0.85 else 相关 if similarity 0.5 else 不相关 }) df pd.DataFrame(result_data) st.dataframe(df, use_container_widthTrue) # 显示详细结果 st.subheader(详细分析) for result in st.session_state.results: similarity result[similarity] if similarity 0.85: st.success(f**排名第{result[rank]}** - 相似度: {similarity:.4f}) elif similarity 0.5: st.warning(f**排名第{result[rank]}** - 相似度: {similarity:.4f}) else: st.error(f**排名第{result[rank]}** - 相似度: {similarity:.4f}) st.write(f候选句子: {result[sentence]}) st.progress(similarity) st.write(---) else: st.info(请在左侧输入查询内容和候选句子然后点击检索按钮) # 侧边栏信息 with st.sidebar: st.header(关于本工具) st.write( 本工具基于阿里达摩院的StructBERT模型开发 支持将单个查询句子与多个候选句子进行语义相似度比较。 **适用场景** - 知识库检索 - FAQ匹配 - 内容去重 - 语义搜索 ) if st.button(清除结果): if results in st.session_state: del st.session_state.results if query in st.session_state: del st.session_state.query st.rerun()5. 实际应用案例5.1 案例一FAQ问答匹配假设我们有一个电商客服的FAQ库faq_library [ 如何退货, 退货流程是什么, 怎样申请退款, 商品有质量问题怎么办, 快递多久能到, 发货时间是多长, 如何联系客服, 客服电话是多少, 怎么修改订单地址, 能否更改收货地址 ] # 用户查询 user_query 我想把买的东西退掉该怎么操作 # 执行检索 results optimized_batch_search(user_query, faq_library, tokenizer, model, top_k3)这样就能快速找到最相关的FAQ答案。5.2 案例二文档内容检索如果你有很多文档段落想要查找相关内容document_paragraphs [ 人工智能是计算机科学的一个分支致力于创建能够执行通常需要人类智能的任务的系统。, 机器学习是人工智能的一个子集它使计算机能够在没有明确编程的情况下学习和做出决策。, 深度学习是机器学习的一个特定子领域它使用多层神经网络来处理复杂的数据模式。, 自然语言处理是人工智能的一个领域专注于计算机与人类语言之间的交互。, 计算机视觉使计算机能够从图像和视频中识别和理解视觉信息。 ] # 查找与神经网络相关的内容 search_query 关于神经网络的技术 results optimized_batch_search(search_query, document_paragraphs, tokenizer, model)6. 性能优化建议6.1 预处理和缓存对于固定的候选句子库可以预先计算好向量并缓存class SentenceDatabase: def __init__(self, tokenizer, model): self.tokenizer tokenizer self.model model self.sentences [] self.embeddings None def add_sentences(self, new_sentences): 添加新句子并更新向量数据库 self.sentences.extend(new_sentences) # 批量计算新句子的向量 new_embeddings [] for sentence in new_sentences: embedding get_sentence_embedding(sentence, self.tokenizer, self.model) new_embeddings.append(embedding.cpu()) # 移到CPU内存保存 if self.embeddings is None: self.embeddings torch.cat(new_embeddings, dim0) else: self.embeddings torch.cat([self.embeddings] new_embeddings, dim0) def search(self, query_sentence, top_k5): 在数据库中进行搜索 query_embedding get_sentence_embedding(query_sentence, self.tokenizer, self.model) # 计算相似度所有向量都在CPU上 similarities cosine_similarity( query_embedding.cpu(), self.embeddings ) # 获取top_k结果 top_scores, top_indices torch.topk(similarities, kmin(top_k, len(self.sentences))) results [] for score, idx in zip(top_scores, top_indices): results.append({ sentence: self.sentences[idx], similarity: score.item(), rank: 0 }) # 添加排名 for rank, result in enumerate(results, 1): result[rank] rank return results # 使用示例 db SentenceDatabase(tokenizer, model) db.add_sentences(faq_library) # 一次性添加所有句子 # 后续查询非常快速 results db.search(怎么退货, top_k3)6.2 批量处理优化对于大量数据的处理还可以进一步优化def batch_encode_sentences(sentences, tokenizer, model, batch_size32): 批量编码句子提高效率 all_embeddings [] for i in range(0, len(sentences), batch_size): batch sentences[i:ibatch_size] inputs tokenizer( batch, return_tensorspt, paddingTrue, truncationTrue, max_length512 ) with torch.no_grad(): outputs model(**inputs.to(cuda)) # 均值池化 attention_mask inputs[attention_mask] token_embeddings outputs.last_hidden_state input_mask_expanded attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float().to(cuda) sum_embeddings torch.sum(token_embeddings * input_mask_expanded, 1) sum_mask torch.clamp(input_mask_expanded.sum(1), min1e-9) batch_embeddings sum_embeddings / sum_mask all_embeddings.append(batch_embeddings.cpu()) return torch.cat(all_embeddings, dim0)7. 总结通过本次代码改造我们将原本只能处理单句对比的nlp_structbert_sentence-similarity_chinese-large工具升级为了支持批量检索的强大工具。主要改造包括核心算法扩展实现了批量句子相似度计算功能性能优化支持分批处理避免内存溢出界面升级改造Streamlit界面支持批量输入和结果展示实用功能添加了排序、分页、可视化等功能改造后的优势处理能力从一次比较2个句子扩展到处理上千个句子使用场景支持知识库检索、FAQ匹配、内容去重等实际应用用户体验直观的界面和清晰的结果展示性能表现优化后的批量处理大幅提高效率现在你可以使用这个改造后的工具来处理各种实际的中文语义匹配任务了。无论是构建智能客服系统、文档检索工具还是内容推荐引擎这个批量检索功能都能为你提供强大的中文语义理解能力。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。