import random
import time
import os
import sys
from tqdm import tqdm
import numpy as np
import paddle
from paddle import inference
from paddlenlp.transformers import AutoModel, AutoTokenizer
from paddlenlp.data import Stack, Tuple, Pad
from paddlenlp.datasets import load_dataset
from paddlenlp.utils.log import logger
sys.path.append('.')
def convert_example(example,
tokenizer,
max_seq_length=512,
pad_to_max_seq_len=False):
result = []
for key, text in example.items():
encoded_inputs = tokenizer(text=text,
max_seq_len=max_seq_length,
pad_to_max_seq_len=pad_to_max_seq_len)
input_ids = encoded_inputs["input_ids"]
token_type_ids = encoded_inputs["token_type_ids"]
result += [input_ids, token_type_ids]
return result
model_dir='./output/yysy/'
corpus_file='./datasets/yysy/milvus/milvus_data_s.csv'
max_seq_length=64
batch_size=64
device='gpu'
cpu_threads=8
model_name_or_path='rocketqa-zh-base-query-encoder'
class Predictor(object):
def __init__(self,
model_dir,
device="gpu",
max_seq_length=128,
batch_size=32,
use_tensorrt=False,
precision="fp32",
cpu_threads=10,
enable_mkldnn=False):
self.max_seq_length = max_seq_length
self.batch_size = batch_size
model_file = model_dir + "inference.get_pooled_embedding.pdmodel"
params_file = model_dir + "inference.get_pooled_embedding.pdiparams"
if not os.path.exists(model_file):
raise ValueError("not find model file path {}".format(model_file))
if not os.path.exists(params_file):
raise ValueError("not find params file path {}".format(params_file))
config = paddle.inference.Config(model_file, params_file)
if device == "gpu":
config.enable_use_gpu(100, 0)
precision_map = {
"fp16": inference.PrecisionType.Half,
"fp32": inference.PrecisionType.Float32,
"int8": inference.PrecisionType.Int8
}
precision_mode = precision_map[precision]
if use_tensorrt:
config.enable_tensorrt_engine(max_batch_size=batch_size,
min_subgraph_size=30,
precision_mode=precision_mode)
elif device == "cpu":
config.disable_gpu()
if enable_mkldnn:
config.set_mkldnn_cache_capacity(10)
config.enable_mkldnn()
config.set_cpu_math_library_num_threads(cpu_threads)
elif device == "xpu":
config.enable_xpu(100)
config.switch_use_feed_fetch_ops(False)
self.predictor = paddle.inference.create_predictor(config)
self.input_handles = [
self.predictor.get_input_handle(name)
for name in self.predictor.get_input_names()
]
self.output_handle = self.predictor.get_output_handle(
self.predictor.get_output_names()[0])
def predict(self, data, tokenizer):
batchify_fn = lambda samples, fn=Tuple(
Pad(axis=0, pad_val=tokenizer.pad_token_id, dtype="int64"), # input
Pad(axis=0, pad_val=tokenizer.pad_token_type_id, dtype="int64"
), # segment
): fn(samples)
all_embeddings = []
examples = []
for idx, text in enumerate(tqdm(data)):
input_ids, segment_ids = convert_example(
text,
tokenizer,
max_seq_length=self.max_seq_length,
pad_to_max_seq_len=True)
examples.append((input_ids, segment_ids))
if (len(examples) >=self.batch_size):
input_ids, segment_ids = batchify_fn(examples)
self.input_handles[0].copy_from_cpu(input_ids)
self.input_handles[1].copy_from_cpu(segment_ids)
self.predictor.run()
logits = self.output_handle.copy_to_cpu()
all_embeddings.append(logits)
examples = []
if (len(examples) > 0):
input_ids, segment_ids = batchify_fn(examples)
self.input_handles[0].copy_from_cpu(input_ids)
self.input_handles[1].copy_from_cpu(segment_ids)
self.predictor.run()
logits = self.output_handle.copy_to_cpu()
all_embeddings.append(logits)
all_embeddings = np.concatenate(all_embeddings, axis=0)
np.save('yysy_corpus_embedding', all_embeddings)
def read_text(file_path):
file = open(file_path)
id2corpus = {}
for idx, data in enumerate(file.readlines()):
id2corpus[idx] = data.strip()
return id2corpus
predictor = Predictor(model_dir, device, max_seq_length, batch_size)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
id2corpus = read_text(corpus_file)
corpus_list = [{idx: text} for idx, text in id2corpus.items()]#用来构建索引库的文本
predictor.predict(corpus_list, tokenizer)
MILVUS_HOST = '127.0.0.1'
MILVUS_PORT = 19530
data_dim = 256
top_k = 20
collection_name = 'literature_search'
partition_tag = 'partition_1'
embedding_name = 'embeddings'
index_config = {
"index_type": "IVF_FLAT",
"metric_type": "L2",
"params": {
"nlist": 1000
},
}
search_params = {
"metric_type": "L2",
"params": {
"nprobe": top_k
},
}
from pymilvus import (
connections,
utility,
FieldSchema,
CollectionSchema,
DataType,
Collection,
)
fmt = "\n=== {:30} ===\n"
text_max_len = 1000
fields = [
FieldSchema(name="pk",
dtype=DataType.INT64,
is_primary=True,#主键
auto_id=False,#不自动增长
max_length=100),#id
FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=text_max_len),#text
FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=data_dim)#embedding
]
schema = CollectionSchema(fields, "Neural Search Index")
class VecToMilvus():#语义向量-->milvus
def __init__(self):
print(fmt.format("start connecting to Milvus"))
connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT)
self.collection = None
def has_collection(self, collection_name):
try:
has = utility.has_collection(collection_name)
print(f"Does collection {collection_name} exist in Milvus: {has}")
return has
except Exception as e:
print("Milvus has_table error:", e)
def creat_collection(self, collection_name):
try:
print(fmt.format("Create collection {}".format(collection_name)))
self.collection = Collection(collection_name,
schema,
consistency_level="Strong")
except Exception as e:
print("Milvus create collection error:", e)
def drop_collection(self, collection_name):
try:
utility.drop_collection(collection_name)
except Exception as e:
print("Milvus delete collection error:", e)
def create_index(self, index_name):
try:
print(fmt.format("Start Creating index"))
self.collection.create_index(index_name, index_config)
print(fmt.format("Start loading"))
self.collection.load()
except Exception as e:
print("Milvus create index error:", e)
def has_partition(self, partition_tag):
try:
result = self.collection.has_partition(partition_tag)
return result
except Exception as e:
print("Milvus has partition error: ", e)
def create_partition(self, partition_tag):
try:
self.collection.create_partition(partition_tag)
print('create partition {} successfully'.format(partition_tag))
except Exception as e:
print('Milvus create partition error: ', e)
def insert(self, entities, collection_name, index_name, partition_tag=None):
try:
if not self.has_collection(collection_name):
self.creat_collection(collection_name)
self.create_index(index_name)
else:
self.collection = Collection(collection_name)
if (partition_tag
is not None) and (not self.has_partition(partition_tag)):
self.create_partition(partition_tag)
self.collection.insert(entities, partition_name=partition_tag)
print(
f"Number of entities in Milvus: {self.collection.num_entities}"
) # check the num_entites
except Exception as e:
print("Milvus insert error:", e)
class RecallByMilvus():#从milvus召回向量
def __init__(self):
print(fmt.format("start connecting to Milvus"))
connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT)
self.collection = None
def get_collection(self, collection_name):
try:
print(fmt.format("Connect collection {}".format(collection_name)))
self.collection = Collection(collection_name)
except Exception as e:
print("Milvus create collection error:", e)
def search(self,
vectors,
embedding_name,
collection_name,
partition_names=[],
output_fields=[]):
try:
self.get_collection(collection_name)
result = self.collection.search(vectors,
embedding_name,
search_params,
limit=top_k,
partition_names=partition_names,
output_fields=output_fields)
return result
except Exception as e:
print('Milvus recall error: ', e)
data_path='./datasets/yysy/milvus/milvus_data_s.csv'
embedding_path='./yysy_corpus_embedding.npy'
index=18
batch_size=5000
def read_text(file_path):
file = open(file_path)
id2corpus = []
for idx, data in enumerate(file.readlines()):
id2corpus.append(data.strip())
return id2corpus
corpus_list_embed=read_text(data_path)
corpus_list_embed[:5]
embeddings = np.load(embedding_path)
embedding_ids = [i for i in range(embeddings.shape[0])]#嵌入ids
client = VecToMilvus()
client.has_collection(collection_name)
client.drop_collection(collection_name)
data_size = len(embedding_ids)
x=[corpus_list_embed[j][:1000]for j in range(10000, 15000,1)]#[:200]文本切片操作
max([len(i) for i in x])
for i in range(0, data_size, batch_size):
print(i)
for i in range(0, data_size, batch_size):#i:0-5000-10000-....
cur_end = i + batch_size
if (cur_end > data_size):#确保下标不越界
cur_end = data_size
batch_emb = embeddings[np.arange(i, cur_end)]#一个批次的嵌入向量
entities = [
[j for j in range(i, cur_end, 1)],#索引
[corpus_list_embed[j][:text_max_len - 1] for j in range(i, cur_end, 1)],#文本
batch_emb #每个批次嵌入向量
]
client.insert(collection_name=collection_name,
entities=entities,
index_name=embedding_name,
partition_tag=partition_tag)
recall_client = RecallByMilvus()
embeddings = embeddings[np.arange(index, index + 1)]
time_start = time.time()# start
result = recall_client.search(embeddings,
embedding_name,
collection_name,
partition_names=[partition_tag],
output_fields=['pk', 'text'])
time_end = time.time()# end
sum_t = time_end - time_start
print('time cost', sum_t, 's')
for hits in result:
for hit in hits:
print(f"hit: {hit}, text field: {hit.entity.get('text')}")