license: apache-2.0 datasets:

• bigcode/the-stack
• bigcode/the-stack-v2
• bigcode/starcoderdata
• bigcode/commitpack library_name: transformers tags:
• code model-index:
• name: Mellum-4b-base results:
  • task: type: text-generation dataset: type: tianyang/repobench_python_v1.1 name: RepoBench 1.1 (Python) metrics:
    • name: EM type: exact_match value: 0.2591 verified: false
    • name: EM ≤ 8k type: exact_match value: 0.2797 verified: false
  • task: type: text-generation dataset: type: tianyang/repobench_python_v1.1 name: RepoBench 1.1 (Python, 2k) metrics:
    • name: EM type: exact_match value: 0.2820 verified: false
  • task: type: text-generation dataset: type: tianyang/repobench_python_v1.1 name: RepoBench 1.1 (Python, 4k) metrics:
    • name: EM type: exact_match value: 0.2795 verified: false
  • task: type: text-generation dataset: type: tianyang/repobench_python_v1.1 name: RepoBench 1.1 (Python, 8k) metrics:
    • name: EM type: exact_match value: 0.2777 verified: false
  • task: type: text-generation dataset: type: tianyang/repobench_python_v1.1 name: RepoBench 1.1 (Python, 12k) metrics:
    • name: EM type: exact_match value: 0.2453 verified: false
  • task: type: text-generation dataset: type: tianyang/repobench_python_v1.1 name: RepoBench 1.1 (Python, 16k) metrics:
    • name: EM type: exact_match value: 0.2110 verified: false
  • task: type: text-generation dataset: type: tianyang/repobench_java_v1.1 name: RepoBench 1.1 (Java) metrics:
    • name: EM type: exact_match value: 0.2858 verified: false
    • name: EM ≤ 8k type: exact_match value: 0.3108 verified: false
  • task: type: text-generation dataset: type: tianyang/repobench_java_v1.1 name: RepoBench 1.1 (Java, 2k) metrics:
    • name: EM type: exact_match value: 0.3202 verified: false
  • task: type: text-generation dataset: type: tianyang/repobench_java_v1.1 name: RepoBench 1.1 (Java, 4k) metrics:
    • name: EM type: exact_match value: 0.3212 verified: false
  • task: type: text-generation dataset: type: tianyang/repobench_java_v1.1 name: RepoBench 1.1 (Java, 8k) metrics:
    • name: EM type: exact_match value: 0.2910 verified: false
  • task: type: text-generation dataset: type: tianyang/repobench_java_v1.1 name: RepoBench 1.1 (Java, 12k) metrics:
    • name: EM type: exact_match value: 0.2492 verified: false
  • task: type: text-generation dataset: type: tianyang/repobench_java_v1.1 name: RepoBench 1.1 (Java, 16k) metrics:
    • name: EM type: exact_match value: 0.2474 verified: false
  • task: type: text-generation dataset: type: gonglinyuan/safim name: SAFIM metrics:
    • name: pass@1 type: pass@1 value: 0.3811 verified: false
  • task: type: text-generation dataset: type: gonglinyuan/safim name: SAFIM (Algorithmic) metrics:
    • name: pass@1 type: pass@1 value: 0.2530 verified: false
  • task: type: text-generation dataset: type: gonglinyuan/safim name: SAFIM (Control) metrics:
    • name: pass@1 type: pass@1 value: 0.3839 verified: false
  • task: type: text-generation dataset: type: gonglinyuan/safim name: SAFIM (API) metrics:
    • name: pass@1 type: pass@1 value: 0.5065 verified: false
  • task: type: text-generation dataset: type: loubnabnl/humaneval_infilling name: HumanEval Infilling (Single-Line) metrics:
    • name: pass@1 type: pass@1 value: 0.6621 verified: false
  • task: type: text-generation dataset: type: loubnabnl/humaneval_infilling name: HumanEval Infilling (Multi-Line) metrics:
    • name: pass@1 type: pass@1 value: 0.3852 verified: false
  • task: type: text-generation dataset: type: loubnabnl/humaneval_infilling name: HumanEval Infilling (Random Span) metrics:
    • name: pass@1 type: pass@1 value: 0.2969 verified: false

モデル概要

Mellum-4b-baseは、JetBrainsが初めて公開したコード関連タスクに最適化された大規模言語モデル(LLM)です。

複数のプログラミング言語にわたる8192トークンのコンテキストウィンドウで4兆トークン以上を学習しており、コード補完に特化しています。 このモデルは40億パラメータのLLaMAスタイルのアーキテクチャを採用しており、クラウド推論(vLLM経由)とローカル展開(llama.cppやOllama使用)の両方に効率的です。

Mellumは自動混合精度(AMP)とbf16精度を使用して学習されました。 Hugging Faceにアップロードされたバージョンは、公開利用のためにbf16形式を保持しています。

プロフェッショナルな開発者ツール(IDEでのインテリジェントなコード提案など)、AI駆動のコーディングアシスタント、コード理解と生成の研究への統合を目的として設計されており、教育用途やファインチューニング実験にも適しています。

このリリースにはベースモデルとPython SFTモデルが含まれています。 他の言語のモデルも間もなく公開予定です。 ベースモデルはすぐに下流タスクに使用できるようにはファインチューニングされていませんが、特定のアプリケーションに適応させるための教師ありファインチューニング(SFT)と強化学習(RL)を完全にサポートしています。

学習データ

• 総学習トークン数: ~4.2兆トークン
• コーパス: The Stack、StarCoderトレーニングデータセット、The Stack v2、CommitPack、英語版Wikipedia

学習詳細

• コンテキストウィンドウ: 8,192トークン
• 最適化: 標準的な言語モデリング目的関数
• ハードウェア: Infiniband接続の256台のH200 NVIDIA GPUクラスタ
• 学習期間: ~20日間

ベンチマーク

ベースモデルのスコアに加えて、Python用にファインチューニングされたMellumのスコアも提供し、モデルユーザーに潜在的な能力の見積もりを提供します。

RepoBench 1.1

• タイプ: 単一行
• 言語: PythonとJava
• 指標: 完全一致(EM)、%

Mellumの最大コンテキストウィンドウは8kであるため、ここでは評価されたすべてのコンテキスト長(2k、4k、8k、12k、16k)にわたる平均性能と、サポート範囲内のコンテキスト長(≤8k)の平均の両方を報告します。

Pythonサブセット

Javaサブセット

構文を考慮した中間挿入(SAFIM)

• タイプ: 複数行と単一行の混合
• 言語: 複数言語
• 指標: pass@1、%

HumanEval中間挿入

• タイプ: 単一行と複数行
• 言語: Python
• 指標: pass@1、%

モデルの改善に取り組んでおり、次のイテレーションを間もなく共有する予定です。

制限事項

• バイアス: 公開コードベースに存在するバイアスを反映する可能性があります。例えば、オープンソースリポジトリと同様のスタイルのコードを生成する傾向があります。
• セキュリティ: コード提案は安全である、または脆弱性がないと想定すべきではありません。

使用例

モデルの実行とサンプリング方法の例を以下に示します。

一般的な生成

from transformers import AutoTokenizer, AutoModelForCausalLM

example = """
import sys
import os
import time

sys.path.append(os.getcwd())

from cluster.prepare_data import get_headers_pairs_list, write_dist_matrix
from cluster.token_edit_distance import get_distance_matrix

if len(sys.argv) < 3:
    print(
        "Too few arguments. You should provide: \n1. dataset_filename" +
        "\n2. output_data_filename"
    )
    sys.exit()

start = time.perf_counter()
dataset_filename_ = sys.argv[1]
output_data_filename_ = sys.argv[2]

headers_pairs = get_headers_pairs_list(dataset_filename_, verbose=True)

dist_matrix, max_dist = get_distance_matrix(
    list(map(lambda x: x[1], headers_pairs)),
    verbose=True
)

write_dist_matrix(dist_matrix, max_dist, output_data_filename_, verbose=True)

end = time.perf_counter()
"""

tokenizer = AutoTokenizer.from_pretrained('JetBrains/Mellum-4b-base')
model = AutoModelForCausalLM.from_pretrained('JetBrains/Mellum-4b-base')
encoded_input = tokenizer(example, return_tensors='pt', return_token_type_ids=False)
input_len = len(encoded_input["input_ids"][0])
out = model.generate(
    **encoded_input,
    max_new_tokens=100,
)
print("### コンテキスト")
print(tokenizer.decode(out[0][:input_len]))
print("### 予測")
print(tokenizer.decode(out[0][input_len:]))

追加ファイルをコンテキストとした中間挿入生成

example = """<filename>utils.py
def multiply(x, y):
    return x * y
<filename>config.py
DEBUG = True
MAX_VALUE = 100
<filename>example.py
<fim_suffix> 

# 関数をテスト
result = calculate_sum(5, 10)
print(result)<fim_prefix>def calculate_sum(a, b):
<fim_middle>"""

encoded_input = tokenizer(example, return_tensors='pt', return_token_type_ids=False)
out = model.generate(
    **encoded_input,
    max_new_tokens=100,
)

引用

このモデルを使用する場合は、以下のように引用してください:

@misc{Mellum-4b-base,
  title     = {Mellum-4b-base},
  author    = {Pavlichenko, Nikita and Nazarov, Iurii and Dolgov, Ivan and Garanina, Ekaterina and Lasocki, Karol and Reshetnikova, Julia and Boitsov, Sergei and Bondyrev, Ivan and Karaeva, Dariia and Sheptyakov, Maksim and Ustalov, Dmitry and Mukhin, Artem and Proshev, Semyon and Abramov, Nikita and Kolomyttseva, Olga and Lysaniuk, Kseniia and Zavidnyi, Ilia and Semenkin, Anton and Tankov, Vladislav and Sazanovich, Uladzislau},
  year      = {2025},
}

連絡先

質問、コラボレーション、リクエストについては、mellum@jetbrains.comまでご連絡ください。