From 83cecc9279e9f0f67f90b3c77d73d438e263c70b Mon Sep 17 00:00:00 2001 From: Meng Zhang Date: Thu, 13 Apr 2023 16:25:13 +0800 Subject: [PATCH] feat: change language field to string to make simplier (#100) --- tabby/server/backend/python.py | 5 ++++- tabby/server/backend/triton.py | 4 +++- tabby/server/models.py | 4 ++-- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/tabby/server/backend/python.py b/tabby/server/backend/python.py index e158dbd..7813c0f 100644 --- a/tabby/server/backend/python.py +++ b/tabby/server/backend/python.py @@ -43,7 +43,10 @@ class PythonModelService: ) def generate(self, request: CompletionRequest) -> List[Choice]: - preset = LanguagePresets[request.language] + preset = LanguagePresets.get(data.language, None) + if preset is None: + return [] + input_ids = self.tokenizer.encode(request.prompt, return_tensors="pt").to( self.device ) diff --git a/tabby/server/backend/triton.py b/tabby/server/backend/triton.py index fb104b6..41aadd5 100644 --- a/tabby/server/backend/triton.py +++ b/tabby/server/backend/triton.py @@ -40,7 +40,9 @@ class TritonService: np_type = np.uint32 model_name = "fastertransformer" - preset = LanguagePresets[data.language] + preset = LanguagePresets.get(data.language, None) + if preset is None: + return [] if self.rewriter: prompt = self.rewriter(preset, data.prompt) diff --git a/tabby/server/models.py b/tabby/server/models.py index 8c68332..a5b9a24 100644 --- a/tabby/server/models.py +++ b/tabby/server/models.py @@ -18,7 +18,7 @@ class Language(str, Enum): class CompletionRequest(BaseModel): - language: Language = Field( + language: str = Field( example=Language.PYTHON, default=Language.UNKNOWN, description="Language for completion request", @@ -48,7 +48,7 @@ class Event(BaseModel): class CompletionEvent(Event): id: str - language: Language + language: str prompt: str created: int choices: List[Choice]