Feature: support code completion (#3612)

python/sglang/srt/code_completion_parser.py

@@ -0,0 +1,174 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Completion templates."""
+
+
+import dataclasses
+import json
+import logging
+import os
+from enum import auto
+
+from sglang.srt.openai_api.protocol import ChatCompletionRequest
+
+logger = logging.getLogger(__name__)
+completion_template_name = None
+
+
+class FimPosition:
+    """Postion of fim middle token."""
+
+    MIDDLE = auto()
+    END = auto()
+
+
+@dataclasses.dataclass
+class CompletionTemplate:
+    """A class that manages completion prompt templates. only for code completion currently."""
+
+    # The name of this template
+    name: str
+
+    # the fim begin token
+    fim_begin_token: str
+
+    # The fim middle token
+    fim_middle_token: str
+
+    # The fim end token
+    fim_end_token: str
+
+    # The position of the fim middle token
+    fim_position: FimPosition
+
+
+# A global registry for all completion templates
+completion_templates: dict[str, CompletionTemplate] = {}
+
+
+def load_completion_template_for_openai_api(completion_template_arg):
+    global completion_template_name
+
+    logger.info(
+        f"Use completion template for the OpenAI-compatible API server: {completion_template_arg}"
+    )
+
+    if not completion_template_exists(completion_template_arg):
+        if not os.path.exists(completion_template_arg):
+            raise RuntimeError(
+                f"Completion template {completion_template_arg} is not a built-in template name "
+                "or a valid completion template file path."
+            )
+
+        assert completion_template_arg.endswith(
+            ".json"
+        ), "unrecognized format of completion template file"
+        with open(completion_template_arg, "r") as filep:
+            template = json.load(filep)
+            try:
+                fim_position = FimPosition[template["fim_position"]]
+            except KeyError:
+                raise ValueError(
+                    f"Unknown fim position: {template['fim_position']}"
+                ) from None
+            register_completion_template(
+                CompletionTemplate(
+                    name=template["name"],
+                    fim_begin_token=template["fim_begin_token"],
+                    fim_middle_token=template["fim_middle_token"],
+                    fim_end_token=template["fim_end_token"],
+                    fim_position=fim_position,
+                ),
+                override=True,
+            )
+        completion_template_name = template["name"]
+    else:
+        completion_template_name = completion_template_arg
+
+
+def register_completion_template(template: CompletionTemplate, override: bool = False):
+    """Register a new completion template."""
+    if not override:
+        assert (
+            template.name not in completion_templates
+        ), f"{template.name} has been registered."
+
+    completion_templates[template.name] = template
+
+
+def completion_template_exists(template_name: str) -> bool:
+    return template_name in completion_templates
+
+
+def is_completion_template_defined() -> bool:
+    global completion_template_name
+    return completion_template_name != None
+
+
+def generate_completion_prompt_from_request(request: ChatCompletionRequest) -> str:
+    global completion_template_name
+    if request.suffix == "":
+        return request.prompt
+
+    return generate_completion_prompt(
+        request.prompt, request.suffix, completion_template_name
+    )
+
+
+def generate_completion_prompt(prompt: str, suffix: str, template_name: str) -> str:
+
+    completion_template = completion_templates[template_name]
+    fim_begin_token = completion_template.fim_begin_token
+    fim_middle_token = completion_template.fim_middle_token
+    fim_end_token = completion_template.fim_end_token
+    fim_position = completion_template.fim_position
+
+    if fim_position == FimPosition.MIDDLE:
+        prompt = f"{fim_begin_token}{prompt}{fim_middle_token}{suffix}{fim_end_token}"
+    elif fim_position == FimPosition.END:
+        prompt = f"{fim_begin_token}{prompt}{fim_end_token}{suffix}{fim_middle_token}"
+
+    return prompt
+
+
+register_completion_template(
+    CompletionTemplate(
+        name="deepseek_coder",
+        fim_begin_token="<｜fim▁begin｜>",
+        fim_middle_token="<｜fim▁hole｜>",
+        fim_end_token="<｜fim▁end｜>",
+        fim_position=FimPosition.MIDDLE,
+    )
+)
+
+
+register_completion_template(
+    CompletionTemplate(
+        name="star_coder",
+        fim_begin_token="<fim_prefix>",
+        fim_middle_token="<fim_middle>",
+        fim_end_token="<fim_suffix>",
+        fim_position=FimPosition.END,
+    )
+)
+
+register_completion_template(
+    CompletionTemplate(
+        name="qwen_coder",
+        fim_begin_token="<|fim_prefix|>",
+        fim_middle_token="<|fim_middle|>",
+        fim_end_token="<|fim_suffix|>",
+        fim_position=FimPosition.END,
+    )
+)

python/sglang/srt/entrypoints/engine.py

@@ -36,6 +36,7 @@ setattr(threading, "_register_atexit", lambda *args, **kwargs: None)
 import torch
 import uvloop
 
+from sglang.srt.code_completion_parser import load_completion_template_for_openai_api
 from sglang.srt.managers.data_parallel_controller import (
     run_data_parallel_controller_process,
 )
@@ -538,6 +539,9 @@ def _launch_subprocesses(
             tokenizer_manager, server_args.chat_template, server_args.model_path
         )
 
+    if server_args.completion_template:
+        load_completion_template_for_openai_api(server_args.completion_template)
+
     # Wait for the model to finish loading
     scheduler_infos = []
     for i in range(len(scheduler_pipe_readers)):

python/sglang/srt/openai_api/adapter.py

@@ -33,6 +33,10 @@ except ImportError:
     # outlines.integrations.utils
     from outlines.integrations.utils import convert_json_schema_to_str
 
+from sglang.srt.code_completion_parser import (
+    generate_completion_prompt_from_request,
+    is_completion_template_defined,
+)
 from sglang.srt.conversation import (
     Conversation,
     SeparatorStyle,
@@ -504,7 +508,11 @@ def v1_generate_request(
                 "To compute logprobs of input prompt, please use the native /generate API."
             )
 
-        prompts.append(request.prompt)
+        prompt = request.prompt
+        if is_completion_template_defined():
+            prompt = generate_completion_prompt_from_request(request)
+        prompts.append(prompt)
+
         lora_paths.append(request.lora_path)
         if request.echo and request.logprobs:
             current_logprob_start_len = 0

python/sglang/srt/server_args.py

@@ -56,6 +56,7 @@ class ServerArgs:
     device: Optional[str] = None
     served_model_name: Optional[str] = None
     chat_template: Optional[str] = None
+    completion_template: Optional[str] = None
     is_embedding: bool = False
     revision: Optional[str] = None
 
@@ -456,6 +457,12 @@ class ServerArgs:
             default=ServerArgs.chat_template,
             help="The buliltin chat template name or the path of the chat template file. This is only used for OpenAI-compatible API server.",
         )
+        parser.add_argument(
+            "--completion-template",
+            type=str,
+            default=ServerArgs.completion_template,
+            help="The buliltin completion template name or the path of the completion template file. This is only used for OpenAI-compatible API server. only for code completion currently.",
+        )
         parser.add_argument(
             "--is-embedding",
             action="store_true",

test/srt/run_suite.py

@@ -70,6 +70,7 @@ suites = {
         TestFile("test_vision_chunked_prefill.py", 223),
         TestFile("test_vision_llm.py", 18.4),
         TestFile("test_vision_openai_server.py", 344),
+        TestFile("test_fim_completion.py", 120),
         TestFile("test_w8a8_quantization.py", 46),
         TestFile("test_eval_fp8_accuracy.py", 172),
         TestFile("test_create_kvindices.py", 2),

test/srt/test_fim_completion.py

@@ -0,0 +1,71 @@
+import unittest
+
+import openai
+
+from sglang.srt.hf_transformers_utils import get_tokenizer
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    popen_launch_server,
+)
+
+
+class TestFimCompletion(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "deepseek-ai/deepseek-coder-1.3b-base"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.api_key = "sk-123456"
+        other_args = ["--completion-template", "deepseek_coder"]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            api_key=cls.api_key,
+            other_args=other_args,
+        )
+        cls.base_url += "/v1"
+        cls.tokenizer = get_tokenizer(cls.model)
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def run_fim_completion(self, number_of_completion):
+        client = openai.Client(api_key=self.api_key, base_url=self.base_url)
+        prompt = "function sum(a: number, b: number): number{\n"
+        suffix = "}"
+
+        prompt_input = self.tokenizer.encode(prompt) + self.tokenizer.encode(suffix)
+        num_prompt_tokens = len(prompt_input) + 2
+
+        response = client.completions.create(
+            model=self.model,
+            prompt=prompt,
+            suffix=suffix,
+            temperature=0.3,
+            max_tokens=32,
+            stream=False,
+            n=number_of_completion,
+        )
+
+        print(response)
+        print(len(response.choices))
+        assert len(response.choices) == number_of_completion
+        assert response.id
+        assert response.created
+        assert response.object == "text_completion"
+        assert (
+            response.usage.prompt_tokens == num_prompt_tokens
+        ), f"{response.usage.prompt_tokens} vs {num_prompt_tokens}"
+        assert response.usage.completion_tokens > 0
+        assert response.usage.total_tokens > 0
+
+    def test_fim_completion(self):
+        for number_of_completion in [1, 3]:
+            self.run_fim_completion(number_of_completion)
+
+
+if __name__ == "__main__":
+    unittest.main()