Mellum 4b Base

🚀 Mellum-4b-base

Mellum-4b-base 是 JetBrains 推出的首個針對代碼相關任務優化的開源大語言模型（LLM）。它在多編程語言上進行訓練，適用於代碼補全、專業開發工具集成等多種場景。

🚀 快速開始

Mellum-4b-base 專為代碼相關任務設計，可用於代碼補全、集成到專業開發工具（如 IDE 中的智能代碼建議）、AI 編碼助手以及代碼理解和生成的研究等。以下是使用示例：

通用生成

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("### Context")
print(tokenizer.decode(out[0][:input_len]))
print("### Prediction")
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> 

# Test the function
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,
)

✨ 主要特性

• 專為代碼任務優化：針對代碼補全進行了專門訓練，在多編程語言上表現出色。
• 高效推理與部署：採用 LLaMA 風格架構，擁有 40 億參數，支持雲推理（如通過 vLLM）和本地部署（如使用 llama.cpp 或 Ollama）。
• 支持多種微調：能夠支持監督微調（SFT）和強化學習（RL），以適應特定應用。

📦 安裝指南

文檔未提供具體安裝步驟，故跳過此章節。

💻 使用示例

基礎用法

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("### Context")
print(tokenizer.decode(out[0][:input_len]))
print("### Prediction")
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> 

# Test the function
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,
)

📚 詳細文檔

模型描述

Mellum-4b-base 是 JetBrains 的首個開源大語言模型，針對代碼相關任務進行了優化。它在超過 4 萬億個標記上進行訓練，上下文窗口為 8192 個標記，適用於多種編程語言。該模型採用 LLaMA 風格架構，擁有 40 億參數，既適合雲推理，也支持本地部署。模型使用自動混合精度（AMP）和 bf16 精度進行訓練，上傳到 Hugging Face 的版本保留了 bf16 格式供公眾使用。此版本包括基礎模型和 Python SFT 模型，其他語言的模型將很快發佈。需要注意的是，基礎模型未針對下游任務進行開箱即用的微調，但完全能夠支持監督微調（SFT）和強化學習（RL）以適應特定應用。

訓練數據

訓練詳情

基準測試

除了基礎模型的得分，還提供了針對 Python 進行微調後的 Mellum 模型的得分，以便用戶對模型的潛在能力有一定的瞭解。

RepoBench 1.1

• 類型：單行
• 語言：Python 和 Java
• 指標：精確匹配（EM），%

由於 Mellum 的最大上下文窗口為 8k，因此報告了所有評估上下文長度（2k、4k、8k、12k 和 16k）的平均性能以及其支持範圍內（≤ 8k）上下文長度的平均性能。

Python 子集

Java 子集

Syntax-Aware Fill-in-the-Middle (SAFIM)

• 類型：多行和單行混合
• 語言：多語言
• 指標：pass@1，%

HumanEval Infilling

• 類型：單行和多行
• 語言：Python
• 指標：pass@1，%

團隊將繼續改進模型，並儘快分享下一個版本。

侷限性

• 偏差：可能反映公共代碼庫中存在的偏差，例如可能生成與開源倉庫風格相似的代碼。
• 安全：代碼建議不能被假定為安全或無漏洞。

🔧 技術細節

文檔未提供具體技術實現細節，故跳過此章節。

📄 許可證

本模型採用 Apache-2.0 許可證。

引用

如果使用此模型，請引用：

@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 聯繫我們。