30 lines
927 B
Python
30 lines
927 B
Python
|
from transformers import AutoTokenizer, AutoModel
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
def load_model(path):
|
|||
|
|
tokenizer = AutoTokenizer.from_pretrained(path)
|
|||
|
|
model = AutoModel.from_pretrained(path)
|
|||
|
|
model.eval()
|
|||
|
|
return tokenizer, model
|
|||
|
|
|
|||
|
|
def embedding(tokenizer,model , sentences):
|
|||
|
|
"""_summary_
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
tokenizer (_type_): 分词器
|
|||
|
|
model (_type_): 向量模型
|
|||
|
|
sentences (_type_): 句子，list
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
_type_: 向量，长度为1024，list
|
|||
|
|
"""
|
|||
|
|
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
|
|||
|
|
with torch.no_grad():
|
|||
|
|
model_output = model(**encoded_input)
|
|||
|
|
# Perform pooling. In this case, cls pooling.
|
|||
|
|
sentence_embeddings = model_output[0][:, 0]
|
|||
|
|
# normalize embeddings
|
|||
|
|
sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
|
|||
|
|
return sentence_embeddings.cpu().numpy().tolist()
|
|||
|
|
|