Rosetta Base Ja

🚀 RoSEtta

RoSEtta (RoFormer-based Sentence Encoder through Distillation) is a general Japanese text embedding model that excels in retrieval tasks. It can handle long sentences with a maximum sequence length of 1024 tokens, runs on a CPU, and is designed to measure semantic similarity between sentences and serve as a retrieval system for passage search based on queries.

✨ Features

• Utilizes RoPE (Rotary Position Embedding).
• Supports a maximum sequence length of 1024 tokens.
• Distilled from large sentence embedding models.
• Specialized for retrieval tasks.

During inference, the prefix "query: " or "passage: " is required. Refer to the Usage section for details.

📚 Documentation

Model Description

This model is based on the RoFormer architecture. After pre-training with MLM loss, weakly supervised learning was performed. Additionally, further training was conducted through distillation using several large embedding models and multi-stage contrastive learning (like GLuCoSE v2).

• Maximum Sequence Length: 1024 tokens
• Output Dimensionality: 768 tokens
• Similarity Function: Cosine Similarity

Usage

Direct Usage (Sentence Transformers)

You can perform inference using SentenceTransformer with the following code:

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

# Download from the 🤗 Hub
# The argument "trust_remote_code=True" is required to load the model
model = SentenceTransformer("pkshatech/RoSEtta-base-ja",trust_remote_code=True)

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

# Get the similarity scores for the embeddings
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]]

Direct Usage (Transformers)

You can perform inference using Transformers with the following code:

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

# Download from the 🤗 Hub
tokenizer = AutoTokenizer.from_pretrained("pkshatech/RoSEtta-base-ja")
# The argument "trust_remote_code=True" is required to load the model
model = AutoModel.from_pretrained("pkshatech/RoSEtta-base-ja",trust_remote_code=True)

# Each input text should start with "query: " or "passage: ".
# For tasks other than retrieval, you can simply use the "query: " prefix.
sentences = [
    'query: PKSHAはどんな会社ですか？',
    'passage: 研究開発したアルゴリズムを、多くの企業のソフトウエア・オペレーションに導入しています。',
    'query: 日本で一番高い山は？',
    'passage: 富士山（ふじさん）は、標高3776.12 m、日本最高峰（剣ヶ峰）の独立峰で、その優美な風貌は日本国外でも日本の象徴として広く知られている。',
]

# Tokenize the input texts
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]

# Get the similarity scores for the embeddings
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]]

Training Details

The fine-tuning of RoSEtta is carried out through the following steps:

1. Pre-training:
   • The model is pre-trained based on the RoFormer architecture.
   • Training data: Japanese Wikipedia and cc100.
2. Weakly supervised learning:
   • Training data: MQA and mc4.
3. Ensemble distillation:
   • The embedded representation was distilled using E5-mistral, gte-Qwen2, and mE5-large as teacher models.
4. Contrastive learning:
   • Triplets were created from JSNLI, MNLI, PAWS-X, JSeM and Mr.TyDi and used for training.
   • This training aimed to improve the overall performance as a sentence embedding model.
5. Search-specific contrastive learning:
   • To make the model more robust for retrieval tasks, additional two-stage training with QA and retrieval tasks was conducted.
   • In the first stage, the synthetic dataset auto-wiki-qa was used for training.
   • In the second stage, JQaRA, MQA, Japanese Wikipedia Human Retrieval, Mr.TyDi,MIRACL, Quiz Works and Quiz No Mor were used.

Benchmarks

Retrieval

Evaluated with MIRACL-ja, JQARA, JaCWIR and MLDR-ja.

Note: Results for OpenAI small embeddings in JQARA and JaCWIR are quoted from JQARA and JaCWIR.

JMTEB

Evaluated with JMTEB. The average score is macro-average.

👨‍💻 Authors

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

📄 License

This model is published under the Apache License, Version 2.0.