add gemma

Author: shell-nlp

gpt_server/model_backend/lmdeploy_backend.py

@@ -7,6 +7,7 @@
 )
 from typing import Any, Dict, AsyncGenerator
 from lmdeploy.archs import get_task
+# from lmdeploy.serve.openai.reasoning_parser import ReasoningParserManager
 from lmdeploy.serve.async_engine import get_names_from_model
 from loguru import logger
 from gpt_server.model_backend.base import ModelBackend
@@ -87,6 +88,8 @@ def __init__(self, model_path) -> None:
         self.messages_type_select = (
             model_type[1] == "base"
         )  # 如果为True 则使用 prompt:str 否则： messages：list
+        # self.reasoning_parser = False
+        # self.tokenizer = self.async_engine.tokenizer
 
     async def stream_chat(self, params: Dict[str, Any]) -> AsyncGenerator:
         prompt = params.get("prompt", "")
@@ -131,12 +134,16 @@ async def stream_chat(self, params: Dict[str, Any]) -> AsyncGenerator:
         results_generator = self.async_engine.generate(
             messages=messages, session_id=int(request_id), gen_config=gen_config
         )
-        text_outputs = ""
+        previous_text = ""
+        current_text = ""
+        previous_token_ids = []
+        current_token_ids = []
+        delta_token_ids = []
         async for request_output in results_generator:
             if await request.is_disconnected():
                 # Abort the request if the client disconnects.
                 await self.async_engine.stop_session(session_id=request_id)
-            text_outputs += request_output.response
+            current_text = current_text + request_output.response
 
             usage = {
                 "prompt_tokens": request_output.input_token_len,
@@ -145,16 +152,39 @@ async def stream_chat(self, params: Dict[str, Any]) -> AsyncGenerator:
                 + request_output.generate_token_len,
             }
             ret = {
-                "text": text_outputs,
+                "text": current_text,
                 "error_code": 0,
                 "usage": usage,
                 "finish_reason": request_output.finish_reason,
             }
+            # if self.reasoning_parser is not None:
+            #     delta_token_ids = (
+            #         request_output.token_ids
+            #         if request_output.token_ids is not None
+            #         else []
+            #     )
+            #     current_token_ids = current_token_ids + delta_token_ids
+            #     reasoning_parser = ReasoningParserManager.get("deepseek-r1")(
+            #         self.tokenizer
+            #     )
+            #     reasoning_delta = reasoning_parser.reasoning_parser.extract_reasoning_content_streaming(
+            #         previous_text=previous_text,
+            #         current_text=current_text,
+            #         delta_text=request_output.response,
+            #         previous_token_ids=previous_token_ids,
+            #         current_token_ids=current_token_ids,
+            #         delta_token_ids=delta_token_ids,
+            #     )
+            #     if reasoning_delta is not None:
+            #         ret["text"] = reasoning_delta.content
+            #         ret["reasoning_content"] = reasoning_delta.reasoning_content
+            #     previous_text = current_text
+            #     previous_token_ids = current_token_ids
             # TODO -------------------------------------------------------------------
             output_info_list = []
             for stop_str in list(stop):
                 if stop_str:
-                    text, bool_value = is_stop(output=text_outputs, stop_str=stop_str)
+                    text, bool_value = is_stop(output=current_text, stop_str=stop_str)
                     output_info_list.append(
                         {"text": text, "bool_value": bool_value, "text_len": len(text)}
                     )
@@ -167,5 +197,5 @@ async def stream_chat(self, params: Dict[str, Any]) -> AsyncGenerator:
                 break
             # TODO -------------------------------------------------------------------
             yield ret
-        logger.info(text_outputs)
+        logger.info(current_text)
         logger.info(usage)

gpt_server/model_worker/base/model_worker_base.py

@@ -53,9 +53,13 @@ def __init__(
         model_type: str = "AutoModel",
         multimodal: bool = False,
     ):
-        self.model_config = AutoConfig.from_pretrained(
-            model_path, trust_remote_code=True
-        )
+        try:
+            self.model_config = AutoConfig.from_pretrained(
+                model_path, trust_remote_code=True
+            )
+        except ValueError as e:
+            logger.warning(e)
+            self.model_config = {}
         # logger.info(f"模型配置：{self.model_config}")
         self.vision_config = getattr(self.model_config, "vision_config", None)
         is_vision = self.vision_config is not None

gpt_server/model_worker/gemma.py

@@ -0,0 +1,83 @@
+import json
+from typing import List
+from fastchat.constants import ErrorCode, SERVER_ERROR_MSG
+import torch
+from loguru import logger
+from gpt_server.model_worker.base.model_worker_base import ModelWorkerBase
+
+
+class GemmaWorker(ModelWorkerBase):
+    def __init__(
+        self,
+        controller_addr: str,
+        worker_addr: str,
+        worker_id: str,
+        model_path: str,
+        model_names: List[str],
+        limit_worker_concurrency: int,
+        conv_template: str = None,  # type: ignore
+    ):
+        super().__init__(
+            controller_addr,
+            worker_addr,
+            worker_id,
+            model_path,
+            model_names,
+            limit_worker_concurrency,
+            conv_template,
+            model_type="AutoModelForCausalLM",
+        )
+        self.stop_words_ids = [1, 106]
+        self.stop = [
+            self.tokenizer.decode(skip_word) for skip_word in self.stop_words_ids
+        ]
+        logger.info(f"{model_names[0]} 停用词: {self.stop}")
+
+    async def generate_stream_gate(self, params):
+        self.call_ct += 1
+        logger.info(f"params {params}")
+        logger.info(f"worker_id: {self.worker_id}")
+        try:
+            messages = params["messages"]
+            if isinstance(messages, list):
+                task = "chat"
+            elif isinstance(messages, str):
+                task = "completion"
+            if task == "chat":
+                text = self.tokenizer.apply_chat_template(
+                    conversation=messages,
+                    tokenize=True,
+                    add_generation_prompt=True,
+                )
+            elif task == "completion":
+                text = messages
+
+            input_ids = self.tokenizer([text], return_tensors="pt").input_ids
+            params["messages"] = messages
+            params["prompt"] = text
+            params["stop"].extend(self.stop)
+            params["stop_words_ids"] = self.stop_words_ids
+            params["input_ids"] = input_ids
+
+            async for ret in self.backend.stream_chat(params=params):
+                response = ret["text"]
+
+                yield json.dumps(ret).encode() + b"\0"
+
+        except torch.cuda.OutOfMemoryError as e:
+            ret = {
+                "text": f"{SERVER_ERROR_MSG}\n\n({e})",
+                "error_code": ErrorCode.CUDA_OUT_OF_MEMORY,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except (ValueError, RuntimeError) as e:
+            logger.info(e)
+            ret = {
+                "text": f"{SERVER_ERROR_MSG}\n\n({e})",
+                "error_code": ErrorCode.INTERNAL_ERROR,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+
+if __name__ == "__main__":
+    GemmaWorker.run()