Bert Base 1024 Biencoder 64M Pairs

🚀 bert-base-1024-biencoder-64M-pairs

このモデルは、MosaicMLの1024シーケンス長で事前学習されたBERTに基づく長文コンテキストのバイエンコーダです。このモデルは、文章や段落を768次元の密ベクトル空間にマッピングし、クラスタリングや意味検索などのタスクに使用できます。

🚀 クイックスタート

📦 インストール

モデルと関連スクリプトのダウンロード

git clone https://huggingface.co/shreyansh26/bert-base-1024-biencoder-64M-pairs

💻 使用例

基本的な使用法

import torch
from torch import nn
from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline, AutoModel
from mosaic_bert import BertModel

# pip install triton==2.0.0.dev20221202 --no-deps if using Pytorch 2.0

class AutoModelForSentenceEmbedding(nn.Module):
    def __init__(self, model, tokenizer, normalize=True):
        super(AutoModelForSentenceEmbedding, self).__init__()

        self.model = model.to("cuda")
        self.normalize = normalize
        self.tokenizer = tokenizer

    def forward(self, **kwargs):
        model_output = self.model(**kwargs)
        embeddings = self.mean_pooling(model_output, kwargs['attention_mask'])
        if self.normalize:
            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)

        return embeddings

    def mean_pooling(self, model_output, attention_mask):
        token_embeddings = model_output[0]  # First element of model_output contains all token embeddings
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

model = AutoModel.from_pretrained("<path-to-model>", trust_remote_code=True).to("cuda")
model = AutoModelForSentenceEmbedding(model, tokenizer)
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')

sentences = ["This is an example sentence", "Each sentence is converted"]

encoded_input = tokenizer(sentences, padding=True, truncation=True, max_length=1024, return_tensors='pt').to("cuda")
embeddings = model(**encoded_input)

print(embeddings)
print(embeddings.shape)

📚 ドキュメント

🔧 技術詳細

トレーニング

このモデルは、Sentence Transformersモデルが使用する同じトレーニングセットからランダムにサンプリングされた6400万対の文章/段落でトレーニングされています。トレーニングセットの詳細はこちらを参照してください。

トレーニング（ハイパーパラメータを含む）、推論、およびデータローディングのスクリプトはすべて、このGithubリポジトリにあります。

評価

いくつかの検索ベースのベンチマーク（CQADupstackEnglishRetrieval、DBPedia、MSMARCO、QuoraRetrieval）でモデルを実行し、結果はこちらにあります。

📄 モデル情報