В блоге 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:
Модель
Описание
Ссылка
Mellum-4b-base
Базовая версия модели
Hugging Face
Mellum-4b-base-gguf
Базовая версия в формате GGUF
Hugging Face
Mellum-4b-sft-python
Дообученная Python-версия
Hugging Face
Mellum-4b-sft-python-gguf
Дообученная Python-версия в формате GGUF
Hugging Face
Источник: habr.com
