Fill in the middle with additional files as context generation
import json
from transformers import AutoTokenizer, AutoModelForCausalLM
example = """
<filename>Utils.kt
package utils
fun multiply(x: Int, y: Int): Int {
return x * y
}
<filename>Config.kt
package config
object Config {
const val DEBUG = true
const val MAX_VALUE = 100
}
<filename>Example.kt
<fim_suffix>
fun main() {
val result = calculateSum(5, 10)
println(result)
}
<fim_prefix>fun calculateSum(a: Int, b: Int): Int {
<fim_middle>
"""
tokenizer = AutoTokenizer.from_pretrained('JetBrains/Mellum-4b-dpo-all')
model = AutoModelForCausalLM.from_pretrained('JetBrains/Mellum-4b-dpo-all')
encoded_input = tokenizer(example, return_tensors='pt', return_token_type_ids=False)
out = model.generate(
**encoded_input,
max_new_tokens=100,
)
Citation
If you use this model, please cite:
@misc{Mellum-4b-dpo-all,
title = {Mellum-4b-dpo-all},
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},
}
For questions, collaborations and requests reach us out via mellum@jetbrains.com