模型卡片：NoInstruct小型嵌入模型v0

NoInstruct嵌入：非对称池化即所需

本模型相比avsolatorio/GIST-small-Embedding-v0具有更优的检索性能。

GIST模型家族在检索任务表现上存在不足。我们提出了一种新方法，在保持对任意指令编码独立性的同时（这是当前检索任务嵌入模型的流行范式），显著提升了检索性能。

模型技术细节将很快发布。

使用方式

from typing import Union
import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

# 加载预训练模型和分词器
model = AutoModel.from_pretrained("avsolatorio/NoInstruct-small-Embedding-v0")
tokenizer = AutoTokenizer.from_pretrained("avsolatorio/NoInstruct-small-Embedding-v0")


def get_embedding(text: Union[str, list[str]], mode: str = "sentence"):
    model.eval()

    assert mode in ("query", "sentence"), f"传入mode={mode}，但仅支持`query`和`sentence`两种模式"

    if isinstance(text, str):
        text = [text]

    # 文本编码
    inp = tokenizer(text, return_tensors="pt", padding=True, truncation=True)

    with torch.no_grad():
        output = model(**inp)

    # 模型优化策略：
    # 查询使用均值池化
    # 句子/文档嵌入使用[CLS]表示

    if mode == "query":
        vectors = output.last_hidden_state * inp["attention_mask"].unsqueeze(2)
        vectors = vectors.sum(dim=1) / inp["attention_mask"].sum(dim=-1).view(-1, 1)
    else:
        vectors = output.last_hidden_state[:, 0, :]

    return vectors


# 示例文本
texts = [
    "REaLTabFormer模型示意图。左侧模块展示使用GPT-2因果语言模型头的非关系型表格数据建模，右侧模块展示如何用序列到序列(Seq2Seq)模型处理关系型数据集的子表。Seq2Seq模型利用父表观测值作为生成子表观测值的条件，且父表训练好的GPT-2模型(权重冻结)也作为Seq2Seq模型的编码器。",
    "人类移动预测具有重要实用价值，应用范围涵盖从增强灾害风险规划到模拟疫情传播。本文提出GeoFormer——一个基于GPT架构改进的解码器专用transformer模型，用于预测人类移动模式。",
    "随着东南亚经济体持续采用数字技术，政策制定者日益关注如何为劳动力应对新兴需求做准备。然而，关于劳动者需要哪些技能来适应这些变化，目前知之甚少。"
]

# 计算嵌入向量
embeddings = get_embedding(texts, mode="sentence")

# 计算句子间余弦相似度
scores = F.cosine_similarity(embeddings.unsqueeze(1), embeddings.unsqueeze(0), dim=-1)
print(scores.cpu().numpy())

# 测试检索性能
query = get_embedding("哪句话讨论了工作相关概念？", mode="query")

scores = F.cosine_similarity(query, embeddings, dim=-1)
print(scores.cpu().numpy())

对Sentence Transformers库的支持即将推出。