JetBrains открыла код Mellum — компактной языковой модели для завершения кода

В блоге JetBrains рассказали о релизе Mellum — компактной языковой модели для завершения кода. Нейросеть не может генерировать код с нуля, но отлично подходит для инструментов, которые помогают его дописывать. У модели всего 4 млрд параметров, поэтому её без проблем можно запустить локально или развернуть на личном сервере.

Mellum поддерживает завершение кода для языков программирования Java, Kotlin, Python, Go, PHP, C, C++, C#, JavaScript, TypeScript, CSS, HTML, Rust и Ruby. У модели относительно небольшое контекстное окно в 8 тыс. токенов. В компании отмечают, что этого достаточно для дополнения кода.

Для кого выпустили Mellum

На Hugging Face инженеры опубликовали базовую модель и дообученную на Python-коде. В компании предлагают использовать нейросети так:

• Запускать в облаке с помощью vLLM, и использовать инференс модели в любых приложениях.

• Запускать локально с помощью llama.cpp или Ollama, что полезно для проектов с конфиденциальной кодовой базой.

• Интегрировать в редакторы кода и среды разработки.

При этом в компании понимают, что мало кто из разработчиков будет использовать Mellum в рабочих задачах. Для этих целей есть более сообразительные модели, которые доступны во всех современных IDE. Нейросеть будет больше полезна исследователям и преподавателям. Например, Mellum может стать практическим пособием по файн-тюнингу базовых языковых моделей.

Как обучали Mellum

Модель обучили приблизительно на 4,2 трлн токенов. В качестве датасета использовали данные:

• The Stack;

• The Stack v2;

• StarCoder Training Dataset;

• CommitPack;

• Англоязычной «Википедии».

В качестве аппаратной платформы использовали кластер из 256 графических ускорителей Nvidia H200, соединённых с помощью коммутаторов InfiniBand. На обучение ушло 20 дней.

Тесты

Mellum сравнили с более крупными языковыми моделями в бенчмарках HumanEval Infilling, RepoBench и SAFIM. Нейросеть ожидаемо уступает им из-за компактности:

Примеры

Вот так можно передать Mellum фрагмент кода на Python и вывести в терминал предложенное нейросетью дополнение:

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:]))

Как получить доступ

Код и веса Mellum доступны на портале Hugging Face: