当前位置: 首页 > news >正文

【人工智能】text2vec-large-chinese模型搭建本地知识库

本demo使用 text2vec-large-chinese 模型进行文本处理,然后再过 bge-reranker-v2-m3进行增强

1. 对文本进行向量处理,并保存只至本地

from sentence_transformers import SentenceTransformer
import torch
import numpy as np
import faiss
import os
import pickle
import fitz  # PyMuPDF
import redevice = torch.device("cuda")
# 加载模型并移动到指定设备
model = SentenceTransformer('/home/sky/model_data/text2vec-large-chinese').to(device)def read_txt(file_path):with open(file_path, 'r', encoding='utf-8') as file:return file.read()def read_pdf(file_path):document = fitz.open(file_path)text = ""for page_num in range(len(document)):page = document.load_page(page_num)text += page.get_text()return textdef preprocess_text(text):# 去除多余的空格和换行符text = re.sub(r'\s+', ' ', text)# 去除特殊字符text = re.sub(r'[^\w\s]', '', text)return text.strip()# 对文本进行分块
def split_into_chunks(text, max_chunk_size=100):# 这里我们简单地以固定长度分块,实际应用中可以根据语义进行更智能的分割words = text.split()chunks = [' '.join(words[i:i + max_chunk_size]) for i in range(0, len(words), max_chunk_size)]return chunks# 读取文档内容
def read_document(file_path):if file_path.endswith('.txt'):return read_txt(file_path)elif file_path.endswith('.pdf'):return read_pdf(file_path)else:raise ValueError("Unsupported file type")# 生成向量
def generate_embeddings(chunks, model, device):embeddings = model.encode(chunks, convert_to_tensor=True, device=device).cpu().numpy()return chunks, embeddings# 保存向量为FAISS格式
def save_embeddings_faiss(chunks, embeddings, output_file):# 创建FAISS索引dimension = embeddings.shape[1]index = faiss.IndexFlatL2(dimension)  # 使用L2距离index.add(embeddings)# 保存FAISS索引faiss.write_index(index, output_file + '.index')# 保存句子列表with open(output_file + '.pkl', 'wb') as f:pickle.dump(chunks, f)# 主函数
def build_knowledge_base(file_path, output_file):# 读取文档内容text = read_document(file_path)# 预处理文本text = preprocess_text(text)# 将文本分割成块chunks = split_into_chunks(text)# 生成向量chunks, embeddings = generate_embeddings(chunks, model, device)# 保存向量为FAISS格式save_embeddings_faiss(chunks, embeddings, output_file)print(f"Embeddings and FAISS index saved to {output_file}.index and {output_file}.pkl")# 使用示例
file_path = '/home/sky/model_data/code/pdf_dir/st.txt'  # 或者 'path/to/your/document.txt'
output_file = '/home/sky/model_data/code/pdf_dir/xldb/textfaiss_new_index'
build_knowledge_base(file_path, output_file)

2. 普通相似度匹配

import faiss
import pickle
import numpy as np
from sentence_transformers import SentenceTransformer
import torch# 检查是否有可用的GPU
device = "cuda"# 加载模型并移动到指定设备
model = SentenceTransformer('/home/sky/model_data/text2vec-large-chinese').to(device)# 加载FAISS索引
index = faiss.read_index('/home/sky/model_data/code/pdf_dir/xldb/textfaiss_index.index')# 加载句子列表
with open('/home/sky/model_data/code/pdf_dir/xldb/textfaiss_index.pkl', 'rb') as f:sentences = pickle.load(f)# 定义查询函数
def search_faiss(query, sentences, index, model, top_n=5):# 生成查询向量query_embedding = model.encode([query], convert_to_tensor=True, device=device).cpu().numpy()# 在FAISS索引中搜索D, I = index.search(query_embedding, top_n)# 返回结果results = []for i, distance in zip(I[0], D[0]):# 1 - L2距离近似于余弦相似度similarity = 1 - distanceresults.append((sentences[i], similarity))return results# 查询示例
query = '今天天气不错。'
results = search_faiss(query, sentences, index, model, top_n=5)# 打印结果
for sentence, score in results:print(f"句子: {sentence},相似度: {score:.4f}")

3. 使用rerank增强 ,rerank模型:bge-reranker-v2-m3

import faiss
import pickle
import numpy as np
from sentence_transformers import SentenceTransformer, CrossEncoder
import torch# 检查是否有可用的GPU
device = "cuda"# 加载模型并移动到指定设备
model = SentenceTransformer('/home/sky/model_data/text2vec-large-chinese').to(device)# 加载FAISS索引
index = faiss.read_index('/home/sky/model_data/code/pdf_dir/xldb/textfaiss_index.index')# 加载句子列表
with open('/home/sky/model_data/code/pdf_dir/xldb/textfaiss_index.pkl', 'rb') as f:sentences = pickle.load(f)# 加载re-ranking模型
reranker = CrossEncoder('/home/sky/model_data/bge-reranker-v2-m3', device=device)# 定义查询函数
def search_faiss(query, sentences, index, model, reranker, top_n=5, rerank_top_n=5):# 生成查询向量query_embedding = model.encode([query], convert_to_tensor=True, device=device).cpu().numpy()# 在FAISS索引中搜索D, I = index.search(query_embedding, rerank_top_n)# 获取初始候选结果initial_results = [(sentences[i], 1 - d) for i, d in zip(I[0], D[0])]# 准备re-ranking的输入rerank_inputs = [(query, sentence) for sentence, _ in initial_results]# 使用re-ranking模型进行重新评分rerank_scores = reranker.predict(rerank_inputs, batch_size=8)# 合并初始分数和re-ranking分数final_results = [(sentence, rerank_score) for (sentence, _), rerank_score in zip(initial_results, rerank_scores)]# 按re-ranking分数排序并取前top_n个结果final_results.sort(key=lambda x: x[1], reverse=True)return final_results[:top_n]# 查询示例
query = '今天天气不错。'
results = search_faiss(query, sentences, index, model, reranker, top_n=5)# 打印结果
for sentence, score in results:print(f"句子: {sentence},相似度: {score:.4f}")

http://www.lryc.cn/news/484545.html

相关文章:

  • 前端入门一之ES6--递归、浅拷贝与深拷贝、正则表达式、es6、解构赋值、箭头函数、剩余参数、String、Set
  • DevOps工程技术价值流:加速业务价值流的落地实践与深度赋能
  • IP数据云 识别和分析tor、proxy等各类型代理
  • vue2 自动化部署 shell 脚本
  • 服务器数据恢复——Ext4文件系统使用fsck后mount不上的数据恢复案例
  • CTF攻防世界小白刷题自学笔记14
  • 家政服务小程序,家政行业数字化发展下的优势
  • Springboot如何打包部署服务器
  • ubuntu将firewall-config导出为.deb文件
  • C++算法练习-day40——617.合并二叉树
  • 2024数维杯国际赛C题【脉冲星定时噪声推断和大气时间信号的时间延迟推断的建模】思路详解
  • 【Linux】MTD 分区
  • MySQL(5)【数据类型 —— 字符串类型】
  • 【数据搜集】初创企业获客,B端数据获取
  • hhdb数据库介绍(9-13)
  • Jmeter基础篇(24)Jmeter目录下有哪些文件夹是可以删除,且不影响使用的呢?
  • 卷积、频域乘积和矩阵向量乘积三种形式之间的等价关系与转换
  • 【Vue】Vue3.0(二十二) v-model 在原始Dom元素、自定义输入组件中双向绑定的底层实现原理详解
  • 史上最强大的 S3 API?介绍 Prompt API。
  • 单片机设计智能翻译手势识别系统
  • 「Mac玩转仓颉内测版12」PTA刷题篇3 - L1-003 个位数统计
  • 飞书文档只读限制复制
  • 【WPF】Prism学习(二)
  • 【鸿蒙开发】第二十一章 Location位置服务
  • 《目标检测》R-CNN网络基础(RCNN,Fast-RCNN)
  • iOS中的定位实现(逆地理编码)及Info.plist位置权限详解
  • 【从零开始的LeetCode-算法】3270. 求出数字答案
  • Web认证机制 Cookie、Token、Session、JWT、OAuth2 解析
  • Docker 基础命令介绍和常见报错解决
  • 如何轻松导出所有 WordPress URL 为纯文本格式