Mellum: Model Cepat untuk Kode di Editormu

Dalam dunia Integrated Development Environments (IDEs), fitur autocompletion telah lama diperkuat oleh kekuatan dan ketepatan. Namun, pendatang baru telah memasuki arena: Artificial Intelligence (AI). Integrasi AI ke dalam IDE semakin umum, mendorong evaluasi ulang alat tradisional. Meskipun metode lama ini masih berfungsi, mereka mungkin tidak lagi memberikan bantuan mutakhir yang diharapkan pengembang saat ini.

Artikel ini diadaptasi dari posting blog oleh JetBrains, dan menyelidiki pelatihan dan kemampuan Mellum, model yang mendasari autocompletion berbasis cloud. JetBrains, yang terkenal dengan IDE dan editor kode yang dirancang untuk programmer, telah mengembangkan Mellum untuk menyediakan autocompletion kode yang cepat dan efisien. Mellum dirancang untuk beroperasi secara efektif pada perangkat lokal, membedakannya sebagai salah satu model tercepat dan paling ringkas di kelasnya. Proyek yang serupa secara konseptual adalah Microsoft Phi.

Menangani Tantangan yang Jelas

Para insinyur JetBrains menghadapi beberapa tantangan penting dalam mengejar autocompletion yang didukung AI:

  • Kecepatan dan Biaya: Model obrolan konvensional terbukti tidak praktis karena biaya komputasi yang tinggi dan waktu respons yang lambat. Model-model ini juga tidak memiliki kesadaran akan teknik khusus kode seperti fill-in-the-middle (FIM) atau token healing.
  • Pemformatan Output: Model obrolan unggulan sering menghasilkan data dalam format yang tidak konsisten, sehingga sulit untuk mengurai respons dan mengintegrasikannya dengan mulus ke dalam editor.
  • Asal Data: Menentukan asal data pelatihan dan mengurangi potensi masalah pelanggaran hak cipta menghadirkan rintangan yang signifikan.

Mellum: Sekilas

Tim pengembangan di JetBrains menyadari bahwa membuat model mereka sendiri adalah pendekatan yang optimal. Tujuan mereka adalah untuk merancang model yang menyeimbangkan kualitas, biaya inferensi, dan latensi sambil dilatih pada data dengan asal yang jelas. Penelitian awal menunjukkan bahwa model dengan sekitar 4 miliar parameter dapat memberikan kemampuan autocompletion yang kompeten untuk berbagai skenario dan pengguna. Selain itu, dengan secara eksklusif melatih model pada kode, mereka dapat membuat kosakata token khusus, bebas dari data yang tidak relevan.

Proses pelatihan untuk model Mellum terdiri dari tiga tahap berbeda, yang masing-masing memberikan pengetahuan baru dan meningkatkan kualitas kode yang dihasilkan. Fase awal melibatkan pra-pelatihan dasar pada korpus besar file individual. Tahap kedua terdiri dari menyempurnakan model dengan sekumpulan contoh khusus yang lebih kecil. Reinforcement Learning with AI Feedback (RLAIF) digunakan pada tahap ketiga untuk menyesuaikan model dengan karakteristik khusus IDE dan menghilangkan output yang tidak diinginkan.

Pra-Pelatihan

Untuk menghindari potensi masalah yang terkait dengan asal data yang ambigu, model dilatih dari bawah, membutuhkan pengantar komprehensif untuk berbagai bahasa, sintaks pemrograman, pola, dan konsep inti.

Dataset

Sumber data utama untuk pra-pelatihan adalah TheStack. Dataset ini memastikan bahwa data tersebut sah secara hukum dan bermanfaat secara praktis.

Proses Pra-Pelatihan

Selama pra-pelatihan, dataset diambil sampelnya beberapa kali untuk mencapai sekitar 3 triliun token. Jendela konteks 8192 token digunakan, dengan dataset dibagi menjadi fragmen dengan ukuran ini. Transformasi fill-in-the-middle (FIM) diterapkan ke setengah dari file di setiap fragmen, mendorong model untuk mempertimbangkan kode sebelumnya dan berikutnya. Teknik ini sangat meniru skenario pembuatan kode dunia nyata.

Fase pra-pelatihan dilakukan pada cluster enam belas node, masing-masing dilengkapi dengan delapan GPU H100. Tahap ini membutuhkan waktu sekitar 15 hari untuk diselesaikan, menghasilkan model Mellum-base dengan 4 miliar parameter.

Pra-pelatihan menciptakan model autocompletion kode tujuan umum dengan pengetahuan luas tentang banyak bahasa pemrograman. Namun, pada tahap ini, model hanya dilatih untuk memprediksi token berikutnya dalam segmen file yang dipilih secara acak. Tanpa konteks tambahan, model tidak memiliki kesadaran akan struktur kode dan tidak memiliki mekanisme untuk menentukan kapan harus berhenti menghasilkan kode.

Tahap fine-tuning dirancang untuk mengatasi keterbatasan ini.

Fine-Tuning Sadar Konteks

Enhanced Fill-in-the-Middle

Tidak seperti pra-pelatihan, di mana fragmen kode dipilih secara acak untuk diprediksi, fine-tuning berkonsentrasi pada segmentasi kode dengan cara yang lebih bermakna, mengajarkan model untuk mengekstrak fragmen kode yang terjadi ‘di alam liar’.

Contoh Khusus

Dalam praktiknya, autocompletion kode memerlukan pemahaman file di sekitarnya dan konteks yang lebih luas, mungkin mencakup seluruh proyek.

Untuk praproses data, perusahaan meluncurkan proyek internal dengan nama kode Code Engine: SDK lintas platform dan serangkaian utilitas konsol yang dikembangkan untuk membangun konteks langsung dari file biasa tanpa memerlukan pengindeksan proyek penuh. SDK ini digunakan di cluster MapReduce internal dan digunakan untuk memproses ribuan repositori publik, menghasilkan banyak contoh berguna untuk pelatihan dalam jangka waktu yang wajar.

Menemukan algoritma yang benar membutuhkan beberapa percobaan dan kesalahan.

Tuning untuk Bahasa Tertentu

Model kecil dapat memperoleh manfaat besar dari spesialisasi untuk bahasa tertentu. Sementara model dasar dilatih pada lebih dari 80 bahasa, sebagian besar pengguna biasanya hanya bekerja dengan satu atau dua. Untuk mengatasi hal ini, JetBrains membuat beberapa model khusus:

  • mellum-all: Mendukung sebagian besar bahasa dan dialek yang tersedia di IDE JetBrains, tetapi kualitas autocompletion lebih rendah daripada model khusus.
  • mellum-python: Mengkhususkan diri dalam Python dan Jupyter.
  • mellum-kotlin: Mengkhususkan diri dalam Java dan Kotlin.
  • mellum-web: Mengkhususkan diri dalam teknologi web.

Langkah Terakhir: RLAIF

Akhirnya, kasus di mana tujuan pembelajaran tidak sesuai dengan harapan pengguna harus diselesaikan. Fase pelatihan tambahan, RLAIF — Reinforcement Learning with AI Feedback, digunakan untuk memecahkan masalah tersebut.

Model belajar dari interaksi pengguna dan memahami cara mencerminkan preferensi pengguna dengan lebih baik.

Pendekatan ini tidak hanya meningkatkan skor kualitas keseluruhan tetapi juga mengurangi jumlah artefak generasi yang mengganggu.

Seberapa Bagus Mellum?

Model ini berkinerja sangat baik untuk ukurannya. Berikut cara evaluasinya:

  • Pertama, model dievaluasi pada tolok ukur internal dengan nama kode ‘JetBrains BigCode’.
  • Kemudian diuji pada tolok ukur publik yang terkenal seperti SAFIM.
  • Akhirnya, statistik penggunaan untuk fitur dikumpulkan, dan metrik pengguna dihitung.

Evaluasi Offline

Mengumpulkan data adalah tugas yang kompleks, tetapi membuat metrik yang baik yang membandingkan saran asli dengan yang baru yang diusulkan oleh jaringan saraf bahkan lebih menantang. Kami melakukan penelitian kecil dan akhirnya menetapkan kombinasi dua metrik utama:

EM:

  • Exact Match adalah ide yang sangat populer.
  • Prediksi dianggap baik jika baris pertama penyelesaian cocok dengan baris pertama dari aslinya, dengan praproses minimal.

KK:

  • Metrik tersebut dinamai menurut nama penulisnya.
  • Jumlah baris yang diusulkan dari aslinya dibagi dengan jumlah baris dalam penyelesaian yang diusulkan.

JetBrains BigCode

Model dievaluasi terhadap dataset tolok ukur yang diperoleh menggunakan alat JetBrains BigCode internal.

Dengan mempertahankan kendali penuh atas dataset kami daripada bergantung pada tolok ukur publik, menjadi mungkin untuk mengevaluasi kualitas model secara andal untuk berbagai gaya dan praktik pengkodean.

Hasil evaluasi JetBrains BigCode kami menunjukkan kualitas yang setara dengan model populer, tetapi Mellum lebih kecil dan lebih efisien.

Kualitas saran baris tunggal (metrik EM)

Tolok Ukur Publik

Model dievaluasi tidak hanya pada dataset internal tetapi juga pada berbagai tolok ukur publik, seperti tolok ukur multibahasa SAFIM (syntax-aware fill in the middle).

Evaluasi Online

Metrik utama disebut rasio kode yang diselesaikan (RoCC). Ini didefinisikan sebagai rasio karakter kode yang ditulis menggunakan autocompletion kode dengan jumlah total kode di editor.

Metrik penting lainnya adalah tingkat penerimaan (AR), yang dihitung sebagai jumlah saran yang diterima dibagi dengan jumlah semua saran yang ditampilkan.

Ini adalah perjalanan yang kompleks, tetapi spesialis JetBrains menyelesaikannya dengan bermartabat. Pada akhirnya, satu model umum dan beberapa model khusus diperoleh, yang tersedia melalui platform JetBrains AI. Mereka sekarang berhasil bekerja di JetBrains AI Assistant.

Apa selanjutnya?

  • Para insinyur JetBrains saat ini sedang mengerjakan model untuk bahasa pengembangan web. Ini mungkin tersedia untuk umum dalam waktu dekat.
  • Ada rencana untuk secara bersamaan meningkatkan jumlah parameter dan keragaman data. Ada banyak tugas berbeda dalam pengkodean — Mellum juga akan dapat melakukannya. Kinerja layanan masih merupakan metrik utama, jadi perluasan model akan berada dalam batas yang wajar.