inf-wse-v2-base-zh

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

queries = ['电脑一体机由什么构成？', '什么是掌上电脑？']
documents = [
    '电脑一体机，是由一台显示器、一个电脑键盘和一个鼠标组成的电脑。',
    '掌上电脑是一种运行在嵌入式操作系统和内嵌式应用软件之上的、小巧、轻便、易带、实用、价廉的手持式计算设备。',
]
input_texts = queries + documents

tokenizer = AutoTokenizer.from_pretrained("infly/inf-wse-v2-base-zh", trust_remote_code=True, use_fast=False)  # Fast tokenizer has not been supported yet
model = AutoModelForMaskedLM.from_pretrained("infly/inf-wse-v2-base-zh", trust_remote_code=True)
model.eval()

max_length = 512

input_batch = tokenizer(input_texts, padding=True, max_length=max_length, truncation=True, return_tensors="pt")
with torch.no_grad():
    embeddings = model(input_batch['input_ids'], input_batch['attention_mask'], return_sparse=False)  # if return_sparse=True, return sparse tensor, else return dense tensor

scores = embeddings[:2] @ embeddings[2:].T
print(scores.tolist())
# [[25.137710571289062, 9.891149520874023], [11.703001976013184, 30.97362518310547]]