🚀 RoSEtta

RoSEtta（基于RoFormer的通过蒸馏得到的句子编码器）是一个通用的日语文本嵌入模型，在检索任务中表现出色。它的最大序列长度为1024，允许输入长句子。该模型可以在CPU上运行，旨在测量句子之间的语义相似度，也可作为基于查询搜索段落的检索系统。

🚀 快速开始

RoSEtta是一个强大的日语文本嵌入模型，可用于测量句子语义相似度和进行段落检索。在推理时，输入文本需以“query: ”或“passage: ”开头。以下是使用示例：

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

# 从🤗 Hub下载
# 加载模型需要参数 "trust_remote_code=True"
model = SentenceTransformer("pkshatech/RoSEtta-base-ja", trust_remote_code=True)

# 每个输入文本应从 "query: " 或 "passage: " 开始。
# 对于检索以外的任务，你可以简单地使用 "query: " 前缀。
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.5910, 0.4332, 0.5421],
# [0.5910, 1.0000, 0.4977, 0.6969],
# [0.4332, 0.4977, 1.0000, 0.7475],
# [0.5421, 0.6969, 0.7475, 1.0000]]

✨ 主要特性

• 使用RoPE（旋转位置嵌入）
• 最大序列长度为1024个标记
• 从大型句子嵌入模型中蒸馏而来
• 专门用于检索任务

📚 详细文档

模型描述

该模型基于RoFormer架构。使用MLM损失进行预训练后，进行了弱监督学习。此外，还通过使用多个大型嵌入模型进行蒸馏和多阶段对比学习（如GLuCoSE v2）进行了进一步训练。

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

使用方法

直接使用（Sentence Transformers）

你可以使用以下代码通过SentenceTransformer进行推理：

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

# 从🤗 Hub下载
# 加载模型需要参数 "trust_remote_code=True"
model = SentenceTransformer("pkshatech/RoSEtta-base-ja", trust_remote_code=True)

# 每个输入文本应从 "query: " 或 "passage: " 开始。
# 对于检索以外的任务，你可以简单地使用 "query: " 前缀。
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.5910, 0.4332, 0.5421],
# [0.5910, 1.0000, 0.4977, 0.6969],
# [0.4332, 0.4977, 1.0000, 0.7475],
# [0.5421, 0.6969, 0.7475, 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

# 从🤗 Hub下载
tokenizer = AutoTokenizer.from_pretrained("pkshatech/RoSEtta-base-ja")
# 加载模型需要参数 "trust_remote_code=True"
model = AutoModel.from_pretrained("pkshatech/RoSEtta-base-ja", trust_remote_code=True)

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

# 对输入文本进行分词
batch_dict = tokenizer(sentences, max_length=1024, 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.5910, 0.4332, 0.5421],
# [0.5910, 1.0000, 0.4977, 0.6969],
# [0.4332, 0.4977, 1.0000, 0.7475],
# [0.5421, 0.6969, 0.7475, 1.0000]]

训练细节

RoSEtta的微调通过以下步骤进行：

1. 预训练：基于RoFormer架构进行预训练。训练数据为日本维基百科和cc100。
2. 弱监督学习：训练数据为MQA和mc4。
3. 集成蒸馏：使用E5-mistral、gte-Qwen2和mE5-large作为教师模型对嵌入表示进行蒸馏。
4. 对比学习：从JSNLI、MNLI、PAWS-X、JSeM和Mr.TyDi创建三元组并用于训练，旨在提高作为句子嵌入模型的整体性能。
5. 特定于搜索的对比学习：为了使模型对检索任务更具鲁棒性，进行了额外的两阶段训练，包括问答和检索任务。第一阶段使用合成数据集auto-wiki-qa进行训练，第二阶段使用JQaRA、MQA、日本维基百科人工检索、Mr.TyDi、MIRACL、Quiz Works和Quiz No Mor进行训练。

基准测试

检索

使用MIRACL-ja、JQARA、JaCWIR和MLDR-ja进行评估。

JMTEB

使用JMTEB进行评估。平均分数为宏平均。

🔧 技术细节

在推理时，输入文本需要添加前缀 "query: " 或 "passage: "，具体使用细节请查看使用部分。

📄 许可证

该模型根据Apache License, Version 2.0发布。

作者

Chihiro Yano, Mocho Go, Hideyuki Tachibana, Hiroto Takegawa, Yotaro Watanabe