🚀 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 联系我们。