Mellum: Mô Hình Nhỏ Gọn Cho IDE

IDE từ lâu đã được hỗ trợ bởi tính năng tự động hoàn thành mạnh mẽ và chính xác. Tuy nhiên, một đối thủ mới đã xuất hiện: Trí tuệ nhân tạo (AI). Việc tích hợp AI vào IDE ngày càng trở nên phổ biến, thúc đẩy việc đánh giá lại các công cụ truyền thống. Mặc dù các phương pháp cũ này vẫn hoạt động, nhưng chúng có thể không còn cung cấp được sự hỗ trợ tiên tiến mà các nhà phát triển hiện nay mong đợi.

Bài viết này được chuyển thể từ một bài đăng trên blog của JetBrains, đi sâu vào quá trình đào tạo và khả năng của Mellum, một mô hình hỗ trợ tính năng tự động hoàn thành dựa trên đám mây. JetBrains, nổi tiếng với các IDE và trình chỉnh sửa code được thiết kế riêng cho các lập trình viên, đã phát triển Mellum để cung cấp khả năng tự động hoàn thành code nhanh chóng và hiệu quả. Mellum được thiết kế để hoạt động hiệu quả trên một thiết bị cục bộ, khiến nó trở thành một trong những mô hình nhanh nhất và nhỏ gọn nhất trong cùng loại. Một dự án tương tự về mặt khái niệm là Microsoft Phi.

Giải Quyết Các Thách Thức Rõ Ràng

Các kỹ sư của JetBrains đã gặp phải một số thách thức quan trọng trong quá trình theo đuổi tính năng tự động hoàn thành dựa trên AI:

  • Tốc độ và Chi phí: Các mô hình trò chuyện thông thường tỏ ra không thực tế do chi phí tính toán cao và thời gian phản hồi chậm. Các mô hình này cũng thiếu nhận thức về các kỹ thuật dành riêng cho code như fill-in-the-middle (FIM) hoặc token healing.
  • Định dạng Đầu ra: Các mô hình trò chuyện hàng đầu thường tạo dữ liệu ở các định dạng không nhất quán, gây khókhăn cho việc phân tích các phản hồi và tích hợp chúng một cách liền mạch vào trình chỉnh sửa.
  • Nguồn gốc Dữ liệu: Xác định nguồn gốc của dữ liệu đào tạo và giảm thiểu các vấn đề tiềm ẩn về vi phạm bản quyền đã đặt ra một trở ngại đáng kể.

Mellum: Tổng Quan

Nhóm phát triển tại JetBrains nhận ra rằng việc tạo ra mô hình của riêng họ là phương pháp tối ưu. Mục tiêu của họ là thiết kế một mô hình cân bằng giữa chất lượng, chi phí suy luận và độ trễ, đồng thời được đào tạo trên dữ liệu có nguồn gốc rõ ràng. Nghiên cứu sơ bộ cho thấy rằng một mô hình với khoảng 4 tỷ tham số có thể cung cấp khả năng tự động hoàn thành có năng lực cho một loạt các tình huống và người dùng. Hơn nữa, bằng cách chỉ đào tạo mô hình trên code, họ có thể thiết lập một từ vựng chuyên biệt gồm các token, không có dữ liệu không liên quan.

Quá trình đào tạo cho mô hình Mellum bao gồm ba giai đoạn riêng biệt, mỗi giai đoạn đóng góp kiến thức mới và nâng cao chất lượng của code được tạo ra. Giai đoạn ban đầu liên quan đến việc tiền đào tạo cơ bản trên một tập hợp lớn các tệp riêng lẻ. Giai đoạn thứ hai bao gồm việc tinh chỉnh mô hình với một tập hợp nhỏ hơn các ví dụ chuyên biệt. Học tăng cường với Phản hồi AI (RLAIF) được sử dụng trong giai đoạn thứ ba để điều chỉnh mô hình theo các đặc điểm cụ thể của IDE và loại bỏ đầu ra không mong muốn.

Tiền Đào Tạo

Để tránh các vấn đề tiềm ẩn liên quan đến nguồn gốc dữ liệu mơ hồ, mô hình đã được đào tạo từ đầu, đòi hỏi một sự giới thiệu toàn diện về nhiều ngôn ngữ, cú pháp lập trình, mẫu và các khái niệm cốt lõi.

Tập Dữ Liệu

Nguồn dữ liệu chính cho tiền đào tạo là TheStack. Tập dữ liệu này đảm bảo rằng dữ liệu vừa hợp pháp vừa có lợi về mặt thực tế.

Quá trình Tiền Đào Tạo

Trong quá trình tiền đào tạo, tập dữ liệu đã được lấy mẫu nhiều lần để đạt khoảng 3 nghìn tỷ token. Một cửa sổ ngữ cảnh gồm 8192 token đã được sử dụng, với tập dữ liệu được chia thành các đoạn có kích thước này. Phép biến đổi fill-in-the-middle (FIM) đã được áp dụng cho một nửa số tệp trong mỗi đoạn, khuyến khích mô hình xem xét cả code trước và sau. Kỹ thuật này mô phỏng chặt chẽ các tình huống tạo code trong thế giới thực.

Giai đoạn tiền đào tạo được thực hiện trên một cụm gồm mười sáu node, mỗi node được trang bị tám GPU H100. Giai đoạn này mất khoảng 15 ngày để hoàn thành, dẫn đến mô hình Mellum-base 4 tỷ tham số.

Tiền đào tạo tạo ra một mô hình tự động hoàn thành code đa năng với kiến thức sâu rộng về nhiều ngôn ngữ lập trình. Tuy nhiên, ở giai đoạn này, mô hình chỉ được đào tạo để dự đoán token tiếp theo trong một phân đoạn tệp được chọn ngẫu nhiên. Nếu không có ngữ cảnh bổ sung, mô hình thiếu nhận thức về cấu trúc code và không có cơ chế để xác định thời điểm dừng tạo code.

Giai đoạn tinh chỉnh được thiết kế để giải quyết những hạn chế này.

Tinh Chỉnh Nhận Biết Ngữ Cảnh

Fill-in-the-Middle Nâng Cao

Không giống như tiền đào tạo, nơi các đoạn code được chọn ngẫu nhiên để dự đoán, tinh chỉnh tập trung vào việc phân đoạn code theo cách có ý nghĩa hơn, dạy mô hình trích xuất các đoạn code xảy ra ‘trong tự nhiên’.

Các Ví dụ Chuyên Biệt

Trong thực tế, tự động hoàn thành code đòi hỏi phải hiểu các tệp xung quanh và các ngữ cảnh rộng hơn, có thể bao gồm toàn bộ dự án.

Đối với quá trình tiền xử lý dữ liệu, công ty đã ra mắt một dự án nội bộ có mật danh Code Engine: một SDK đa nền tảng và một tập hợp các tiện ích dòng lệnh được phát triển để xây dựng ngữ cảnh trực tiếp từ các tệp thông thường mà không yêu cầu lập chỉ mục toàn bộ dự án. SDK này đã được triển khai trên một cụm MapReduce nội bộ và được sử dụng để xử lý hàng nghìn kho lưu trữ công khai, tạo ra nhiều ví dụ hữu ích để đào tạo trong một khung thời gian hợp lý.

Việc tìm ra các thuật toán chính xác đòi hỏi một số thử và sai.

Điều Chỉnh Cho Các Ngôn Ngữ Cụ Thể

Các mô hình nhỏ có thể hưởng lợi rất nhiều từ việc chuyên môn hóa cho các ngôn ngữ cụ thể. Mặc dù mô hình cơ sở được đào tạo trên hơn 80 ngôn ngữ, nhưng hầu hết người dùng thường chỉ làm việc với một hoặc hai ngôn ngữ. Để giải quyết vấn đề này, JetBrains đã tạo ra nhiều mô hình chuyên biệt:

  • mellum-all: Hỗ trợ hầu hết các ngôn ngữ và phương ngữ có sẵn trong IDE của JetBrains, nhưng chất lượng tự động hoàn thành thấp hơn so với các mô hình chuyên biệt.
  • mellum-python: Chuyên về Python và Jupyter.
  • mellum-kotlin: Chuyên về Java và Kotlin.
  • mellum-web: Chuyên về các công nghệ web.

Bước Cuối Cùng: RLAIF

Cuối cùng, các trường hợp mà mục tiêu học tập không phù hợp với mong đợi của người dùng phải được giải quyết. Một giai đoạn đào tạo bổ sung, RLAIF — Học tăng cường với Phản hồi AI, được sử dụng để giải quyết những vấn đề như vậy.

Mô hình học hỏi từ các tương tác của người dùng và hiểu cách phản ánh tốt hơn các tùy chọn của người dùng.

Phương pháp này không chỉ cải thiện điểm chất lượng tổng thể mà còn giảm số lượng các tạo tác tạo khó chịu.

Mellum Tốt Đến Mức Nào?

Mô hình hoạt động đặc biệt tốt so với kích thước của nó. Đây là cách nó được đánh giá:

  • Đầu tiên, mô hình được đánh giá trên một điểm chuẩn nội bộ có mật danh “JetBrains BigCode”.
  • Sau đó, nó được kiểm tra trên các điểm chuẩn công khai nổi tiếng như SAFIM.
  • Cuối cùng, thống kê sử dụng cho các tính năng đã được thu thập và các số liệu người dùng đã được tính toán.

Đánh Giá Ngoại Tuyến

Việc thu thập dữ liệu là một nhiệm vụ phức tạp, nhưng việc tạo ra một số liệu tốt so sánh đề xuất ban đầu với đề xuất mới do mạng nơ-ron đưa ra thậm chí còn khó khăn hơn. Chúng tôi đã thực hiện một nghiên cứu nhỏ và cuối cùng đã quyết định một sự kết hợp của hai số liệu chính:

EM:

  • Exact Match là một ý tưởng rất phổ biến.
  • Một dự đoán được coi là tốt nếu dòng đầu tiên của phần hoàn thành khớp với dòng đầu tiên của bản gốc, với quá trình tiền xử lý tối thiểu.

KK:

  • Số liệu này được đặt theo tên của các tác giả của nó.
  • Số lượng dòng được đề xuất từ bản gốc chia cho số lượng dòng trong phần hoàn thành được đề xuất.

JetBrains BigCode

Mô hình đã được đánh giá dựa trên một tập dữ liệu điểm chuẩn thu được bằng công cụ JetBrains BigCode nội bộ.

Bằng cách duy trì toàn quyền kiểm soát tập dữ liệu của mình thay vì dựa vào các điểm chuẩn công khai, có thể đánh giá một cách đáng tin cậy chất lượng mô hình cho các kiểu và thực hành viết code khác nhau.

Kết quả đánh giá JetBrains BigCode của chúng tôi cho thấy chất lượng tương đương với các mô hình phổ biến, nhưng Mellum nhỏ hơn và hiệu quả hơn.

Chất lượng của các đề xuất một dòng (số liệu EM)

Các Điểm Chuẩn Công Khai

Mô hình đã được đánh giá không chỉ trên tập dữ liệu nội bộ mà còn trên các điểm chuẩn công khai khác nhau, chẳng hạn như điểm chuẩn đa ngôn ngữ SAFIM (syntax-aware fill in the middle).

Đánh Giá Trực Tuyến

Số liệu chính được gọi là tỷ lệ code đã hoàn thành (RoCC). Nó được định nghĩa là tỷ lệ ký tự code được viết bằng cách sử dụng tính năng tự động hoàn thành code so với tổng số lượng code trong trình chỉnh sửa.

Một số liệu quan trọng khác là tỷ lệ chấp nhận (AR), được tính bằng số lượng đề xuất được chấp nhận chia cho số lượng tất cả các đề xuất được hiển thị.

Đây là một hành trình phức tạp, nhưng các chuyên gia của JetBrains đã hoàn thành nó một cách trang trọng. Cuối cùng, một mô hình chung và một số mô hình chuyên biệt đã được thu được, có sẵn thông qua nền tảng JetBrains AI. Chúng hiện đang hoạt động thành công trong JetBrains AI Assistant.

Tiếp theo là gì?

  • Các kỹ sư của JetBrains hiện đang làm việc trên một mô hình cho các ngôn ngữ phát triển web. Nó có thể được cung cấp công khai trong tương lai gần.
  • Có kế hoạch đồng thời tăng cả số lượng tham số và sự đa dạng của dữ liệu. Có rất nhiều nhiệm vụ khác nhau trong viết code — Mellum cũng sẽ có thể thực hiện chúng. Hiệu suất dịch vụ vẫn là một số liệu quan trọng, vì vậy việc mở rộng mô hình sẽ nằm trong giới hạn hợp lý.