GLuCoSE-base-ja-v2 - 专为日本文本检索及句子相似度设计的嵌入模型

GLuCoSE-base-ja-v2项目介绍

项目背景

GLuCoSE-base-ja-v2 是一个专注于日语文本嵌入的模型，尤其在信息检索任务中表现出色。这个模型可以在普通的CPU上运行，以评估语句间的语义相似性，并作为基于查询的段落搜索系统。

项目特点

专为检索任务优化：在类似大小的模型中，其在MIRACL和其他任务中表现最佳。
日语文本优化：针对日语文本的处理进行了特别优化。
CPU可运行：无需复杂的硬件设施，普通的CPU便可以支持模型的运行。

模型描述

GLuCoSE-base-ja-v2 模型基于 GLuCoSE进行构建，并通过大规模嵌入模型的蒸馏与多阶段对比学习进行微调。

最大序列长度：512个标记
输出维度：768个标记
相似性函数：余弦相似性

使用方法

直接使用 (基于 Sentence Transformers)

可以使用 SentenceTransformer 进行推理，代码如下：

from sentence_transformers import SentenceTransformer
import torch.nn.functional as F

# 从 Hugging Face Hub 下载模型
model = SentenceTransformer("pkshatech/GLuCoSE-base-ja-v2")

# 每个输入文本应以 "query: " 或 "passage: " 开头。
sentences = [
    'query: PKSHAはどんな会社ですか？',
    'passage: 研究開発したアルゴリズムを、多くの企業のソフトウエア・オペレーションに導入しています。',
    'query: 日本で一番高い山は？',
    'passage: 富士山（ふじさん）は、標高3776.12 m、日本最高峰（剣ヶ峰）の独立峰で、その優美な風貌は日本国外でも日本の象徴として広く知られている。',
]
embeddings = model.encode(sentences,convert_to_tensor=True)
print(embeddings.shape)
# [4, 768]

# 获取嵌入的相似性得分
similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)
print(similarities)
# [[1.0000, 0.6050, 0.4341, 0.5537],
# [0.6050, 1.0000, 0.5018, 0.6815],
# [0.4341, 0.5018, 1.0000, 0.7534],
# [0.5537, 0.6815, 0.7534, 1.0000]]

直接使用 (基于 Transformers)

还可以使用 Transformers 进行推理，代码如下：

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

def mean_pooling(last_hidden_states: Tensor,attention_mask: Tensor) -> Tensor:
    emb = last_hidden_states * attention_mask.unsqueeze(-1)
    emb = emb.sum(dim=1) / attention_mask.sum(dim=1).unsqueeze(-1)
    return emb

# 从 Hugging Face Hub 下载模型
tokenizer = AutoTokenizer.from_pretrained("pkshatech/GLuCoSE-base-ja-v2")
model = AutoModel.from_pretrained("pkshatech/GLuCoSE-base-ja-v2")

# 每个输入文本应以 "query: " 或 "passage: " 开头。
sentences = [
    'query: PKSHAはどんな会社ですか？',
    'passage: 研究開発したアルゴリズムを、多くの企業のソフトウエア・オペレーションに導入しています。',
    'query: 日本で一番高い山は？',
    'passage: 富士山（ふじさん）は、標高3776.12 m、日本最高峰（剣ヶ峰）の独立峰で、その優美な風貌は日本国外でも日本の象徴として広く知られている。',
]

# 令输入文本进行标记化
batch_dict = tokenizer(sentences, max_length=512, padding=True, truncation=True, return_tensors='pt')

outputs = model(**batch_dict)
embeddings = mean_pooling(outputs.last_hidden_state, batch_dict['attention_mask'])
print(embeddings.shape)
# [4, 768]

# 获取嵌入的相似性得分
similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)
print(similarities)
# [[1.0000, 0.6050, 0.4341, 0.5537],
# [0.6050, 1.0000, 0.5018, 0.6815],
# [0.4341, 0.5018, 1.0000, 0.7534],
# [0.5537, 0.6815, 0.7534, 1.0000]]