infgrad/puff-base-v1

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

# Textos a codificar
texts = ["通用向量编码", "hello world", "支持中英互搜,不建议纯英文场景使用"]
# Directorio del modelo
model_dir = "{MODEL_PATH}"

#### Método 1: Usar SentenceTransformer
# !!!!!!!!!!!!!!El valor por defecto es 4096 dimensiones, si necesitas otras dimensiones, copia los archivos correspondientes de 2_Dense_{dims} a la carpeta 2_Dense!!!!!!!!!!!!!!
model = SentenceTransformer(model_dir)
vectors = model.encode(texts, convert_to_numpy=True, normalize_embeddings=True)
print(vectors.shape)
print(vectors[:, :4])

#### Método 2: Usar la biblioteca transformers
# !!!!!!!!!!!!!! Este código leerá el peso de la capa lineal correspondiente según el valor de vector_dim, ¡elige vector_dim según tus necesidades! !!!!!!!!!!!!!!
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])