import os
import torch
from transformers import AutoModel, AutoTokenizer
from sentence_transformers import SentenceTransformer
from sklearn.preprocessing import normalize

# 待编码文本
texts = ["通用向量编码", "hello world", "支持中英互搜，不建议纯英文场景使用"]
# 模型目录
model_dir = "{MODEL_PATH}"

#### 方法1：使用SentenceTransformer
# ！！！！！！！！！！！！！！默认是4096维度，如需其他维度，请自行复制2_Dense_{dims}中的文件到2_Dense文件夹中覆盖！！！！！！！！！！！！！！
model = SentenceTransformer(model_dir)
vectors = model.encode(texts, convert_to_numpy=True, normalize_embeddings=True)
print(vectors.shape)
print(vectors[:, :4])

#### 方法2：使用transformers库
# ！！！！！！！！！！！！！！ 本代码会根据vector_dim值会读取对应的Linear层权重，请按需选择vector_dim ！！！！！！！！！！！！！！
vector_dim = 4096
model = AutoModel.from_pretrained(model_dir).eval()
tokenizer = AutoTokenizer.from_pretrained(model_dir)
vector_linear = torch.nn.Linear(in_features=model.config.hidden_size, out_features=vector_dim)
vector_linear_dict = {
    k.replace("linear.", ""): v for k, v in
    torch.load(os.path.join(model_dir, f"2_Dense_{vector_dim}/pytorch_model.bin")).items()
}
vector_linear.load_state_dict(vector_linear_dict)
with torch.no_grad():
    input_data = tokenizer(texts, padding="longest", truncation=True, max_length=512, return_tensors="pt")
    attention_mask = input_data["attention_mask"]
    last_hidden_state = model(**input_data)[0]
    last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
    vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
    vectors = normalize(vector_linear(vectors).cpu().numpy())
print(vectors.shape)
print(vectors[:, :4])

infgrad/puff-base-v1

模型介绍

模型目录结构

模型使用方法

infgrad/
puff-base-v1