Codesearch ModernBERT Snake

🚀 Shuu12121/CodeSearch-ModernBERT-Snakeに基づくSentenceTransformer🐍

このモデルは、コード検索に特化したSentence Transformerモデルです。私が一から事前学習を行ったコード特化のModernBERTモデルである Shuu12121/CodeModernBERT-Snake をベースにファインチューニングされています。最大シーケンス長が8192トークンに対応しており、非常に長いコード片やドキュメントにも対応可能です。また、75M程度と極めて小さいモデルながら、コード検索タスクにおいて高い性能を発揮します。

✨ 主な機能

• コード片とドキュメントの意味的類似性を効率的に計算
• 最大シーケンス長が8192トークンに対応
• コンパクトなモデルサイズで高い性能を発揮

📦 インストール

Sentence Transformers をインストールするには、以下のコマンドを実行します。

pip install -U sentence-transformers

💻 使用例

基本的な使用法

from sentence_transformers import SentenceTransformer

# モデルをダウンロードしてロード
model = SentenceTransformer("Shuu12121/CodeSearch-ModernBERT-Snake")

# 推論用の文リスト
sentences = [
    'Encrypts the zip file',
    'def freeze_encrypt(dest_dir, zip_filename, config, opt):\n    \n    pgp_keys = grok_keys(config)\n    icefile_prefix = "aomi-%s" % \\\n                     os.path.basename(os.path.dirname(opt.secretfile))\n    if opt.icefile_prefix:\n        icefile_prefix = opt.icefile_prefix\n\n    timestamp = time.strftime("%H%M%S-%m-%d-%Y",\n                              datetime.datetime.now().timetuple())\n    ice_file = "%s/%s-%s.ice" % (dest_dir, icefile_prefix, timestamp)\n    if not encrypt(zip_filename, ice_file, pgp_keys):\n        raise aomi.exceptions.GPG("Unable to encrypt zipfile")\n\n    return ice_file',
    'def transform(self, sents):\n        \n\n        def convert(tokens):\n            return torch.tensor([self.vocab.stoi[t] for t in tokens], dtype=torch.long)\n\n        if self.vocab is None:\n            raise Exception(\n                "Must run .fit() for .fit_transform() before " "calling .transform()."\n            )\n\n        seqs = sorted([convert(s) for s in sents], key=lambda x: -len(x))\n        X = torch.LongTensor(pad_sequence(seqs, batch_first=True))\n        return X',
]

# 埋め込みベクトルの生成
embeddings = model.encode(sentences)
print(embeddings.shape)  # 出力: [3, 512]

# 類似度スコアの計算
similarities = model.similarity(embeddings, embeddings)
print(similarities.shape)  # 出力: [3, 3]

📚 ドキュメント

モデル評価

本モデルは、75M程度と極めて小さいモデルながら、コードサーチタスクにおける評価指標である CodeSearchNet で 72.12 を達成しました。これは 400Mのパラメータを持つ Salesforce/SFR-Embedding-Code-400M_R と比較しても遜色ないレベルです。他のタスクには対応していないため、評価値は提供されていません。

モデル詳細

ライブラリバージョン

• Python: 3.11.11
• Sentence Transformers: 3.4.1
• Transformers: 4.50.0
• PyTorch: 2.6.0+cu124
• Accelerate: 1.5.2
• Datasets: 3.4.1
• Tokenizers: 0.21.1

📄 ライセンス

このモデルは Apache-2.0 ライセンスの下で提供されています。

📖 引用情報

Sentence Transformers

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084",
}

MultipleNegativesRankingLoss

@misc{henderson2017efficient,
    title={Efficient Natural Language Response Suggestion for Smart Reply},
    author={Matthew Henderson and Rami Al-Rfou and Brian Strope and Yun-hsuan Sung and Laszlo Lukacs and Ruiqi Guo and Sanjiv Kumar and Balint Miklos and Ray Kurzweil},
    year={2017},
    eprint={1705.00652},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}