Skip to content

Add codestral prompt for code completions

Bruno Cardoso requested to merge bc/custom-models-codestral into master

What does this MR do and why?

This MR adds support for using a codestral model for Code Completions via Code Suggestions.

It does so by:

  1. Create a new abstract class AiGatewayCodeCompletionMessage requiring only the prompt method to be implemented by subclasses.
  2. Create the respective CodestralMessage subclass that implements the FIM prompt required by Codestral.
  3. Refactor CodeGemmaMessage to also subclass AiGatewayCodeCompletionMessage, removing some duplicated code.

Part of #467433 (closed)

Evaluation results

I ran the code-suggestions eval pipeline against a local gdk instance that uses codestral for code completions.

config
{
  "beam_config": {
    "pipeline_options": {
      "runner": "DirectRunner",
      "project": "dev-ai-research-0e2f8974",
      "region": "us-central1",
      "temp_location": "gs://prompt-library/tmp/",
      "save_main_session": false
    }
  },
  "input_bq_table": "dev-ai-research-0e2f8974.code_suggestion.input_testcases_v1",
  "output_bq_table": "dev-ai-research-0e2f8974.code_suggestion_experiments.bc_codestral_completions_py_only_temp0",
  "throttle_sec": 1,
  "batch_size": 10,
  "min_length": 0,
  "eval_setups": [
    {
      "llm_config": {
        "name": "gitlab-code-completion",
        "parameters": {
          "base_url": "http://gdk.test:3000",
          "intent": "completion"
        }
      },
      "languages": [
        "python"
      ],
      "post_transformations": [],
      "include_suffix": true
    }
  ]
}
model avg_similarity_score sample_size table
codestral:latest 0.77 510 dev-ai-research-0e2f8974.code_suggestion_experiments.bc_codestral_completions_py_only_temp0
prod (code-gecko) 0.79 510 dev-ai-research-0e2f8974.code_suggestion_experiments.bc_becko_completions_py_only

Example query:

SELECT model, round(avg(similarity_score), 2) as avg_score, count(*) FROM `dev-ai-research-0e2f8974.code_suggestion_experiments.bc_codestral_completions_py_only_temp0` WHERE NOT STARTS_WITH(model_completion, '{') GROUP BY model

MR acceptance checklist

Please evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Screenshots or screen recordings

Screenshots are required for UI changes, and strongly recommended for all other merge requests.

Before After

How to set up and validate locally

Numbered steps to set up and validate the change are strongly suggested.

Edited by Bruno Cardoso

Merge request reports