File size: 1,594 Bytes
974817f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
import torch.distributed
import faiss
import pandas as pd
import faiss
import numpy as np
import jsonlines, json
from transformers import AutoModel
import os
import torch
'''
data format:
{
"bibkey": "some_bibkey",
"text": "The abstract or text of the paper."
}
example:
{
"bibkey": "arxivid1234.5678",
"text": "Title: A Study on Something\nAbstract: This paper discusses the findings of a study on something important in the field of research.\nAuthors: John Doe"
}
'''
model_name = "openbmb/MiniCPM-Embedding-Light"
model = AutoModel.from_pretrained(model_name, trust_remote_code=True, attn_implementation="flash_attention_2", torch_dtype=torch.float16).to("cuda")
input_path = "./data/arxiv.jsonl"
with jsonlines.open(input_path) as f:
survey_data = list(f)
xids = [item["bibkey"] for item in survey_data]
passages = [item["text"] for item in survey_data]
embeddings_doc_dense, _ = model.encode_corpus(passages, max_length=1024)
# faiss save index
index = faiss.IndexFlatIP(embeddings_doc_dense.shape[1])
id_map_index = faiss.IndexIDMap(index)
index = faiss.index_cpu_to_all_gpus(id_map_index)
x_ids_int = np.array(np.arange(len(xids)))
str_int_ids = {}
for i in range(len(xids)):
str_int_ids[xids[i]] = x_ids_int[i]
str_int_ids_df = pd.DataFrame(str_int_ids, index=[0]).T.reset_index()
str_int_ids_df.columns = ["str_id", "int_id"]
str_int_ids_df.to_csv("./index/str_int_ids_abstract.csv", index=False)
index.add_with_ids(embeddings_doc_dense, x_ids_int)
index = faiss.index_gpu_to_cpu(index)
faiss.write_index(index, "./index/index_abstract.faiss")
|