Fix prompt is not string bug (#81)

* Fix prompt is not string bug

* update parameter type

Author: HaiHui886

llmserve/backend/llm/engines/_base.py

@@ -12,7 +12,7 @@
 from llmserve.backend.logger import get_logger
 from llmserve.backend.server.models import Args, Prompt
 import asyncio
-from typing import AsyncGenerator, Generator
+from typing import Union, AsyncGenerator, Generator
 
 logger = get_logger(__name__)
 
@@ -67,5 +67,5 @@ async def check_health(self):
         pass
 
     @abstractmethod
-    def stream_generate_texts(self, prompt: str) -> Generator[str, None, None]:
+    def stream_generate_texts(self, prompt: Union[Prompt, List[Prompt]]) -> Generator[str, None, None]:
         pass

llmserve/backend/llm/engines/generic.py

@@ -32,7 +32,7 @@
 from llmserve.backend.server.utils import render_gradio_params
 from ._base import LLMEngine
 
-from typing import AsyncGenerator, Generator
+from typing import AsyncGenerator, Generator, Union
 from queue import Empty
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
@@ -284,7 +284,7 @@ def ping(self) -> bool:
         """Ping the worker."""
         return True
 
-    async def worker_stream_generate_texts(self, prompt: str, **kwargs) -> Generator[str, None, None]: # type: ignore
+    async def worker_stream_generate_texts(self, prompt: Union[Prompt, List[Prompt]], **kwargs) -> Generator[str, None, None]: # type: ignore
         logger.info(f"Call PredictionWorker.worker_stream_generate_texts with kwargs: {kwargs}")
         for s in self.generator.streamGenerate(prompt, **kwargs):
             # logger.info(f"PredictionWorker.worker_stream_generate_texts -> yield ->{s}")
@@ -430,7 +430,7 @@ async def check_health(self):
                     "Reinitializing worker group."
                 )
 
-    def stream_generate_texts(self, prompt: str) -> Generator[str, None, None]: # type: ignore
+    def stream_generate_texts(self, prompt: Union[Prompt, List[Prompt]]) -> Generator[str, None, None]: # type: ignore
         logger.info(f"GenericEngine.stream_generate_texts -> worker.length: {len(self.base_worker_group)}")
         worker0 = self.base_worker_group[0]
         for strHandle in worker0.worker_stream_generate_texts.remote(

llmserve/backend/llm/pipelines/_base.py

@@ -341,7 +341,7 @@ def _sanitize_parameters(
         return preprocess_params, forward_params, postprocess_params
 
     @abstractmethod
-    def streamGenerate(self, prompt: str, **generate_kwargs) -> Generator[str, None, None]:
+    def streamGenerate(self, prompt: Union[Prompt, List[Prompt]], **generate_kwargs) -> Generator[str, None, None]:
         pass
 
 class StreamingPipeline(BasePipeline):

llmserve/backend/llm/pipelines/default_pipeline.py

@@ -5,7 +5,7 @@
 from transformers import PreTrainedModel, PreTrainedTokenizer
 
 from llmserve.backend.logger import get_logger
-from llmserve.backend.server.models import Response
+from llmserve.backend.server.models import Prompt, Response
 import json
 
 from ._base import BasePipeline
@@ -167,15 +167,22 @@ def postprocess(self, model_outputs, **postprocess_kwargs) -> List[Response]:
             response.postprocessing_time = et
         return decoded
 
-    def streamGenerate(self, prompt: str, **generate_kwargs) -> Generator[str, None, None]:
+    def streamGenerate(self, prompt: Union[Prompt, List[Prompt]], **generate_kwargs) -> Generator[str, None, None]:
         logger.info(f"DefaultPipeline.streamGenerate with generate_kwargs: {generate_kwargs}")
         # timeout=0  will dramatic slow down the speed of generator, the root caused still unknow
         streamer = TextIteratorStreamer(self.tokenizer,
                                         # timeout=0,
                                         skip_prompt=True,
                                         skip_special_tokens=True)
-        input_ids = self.tokenizer([prompt], return_tensors="pt")
-        # generation_kwargs = dict(input_ids, streamer=streamer, max_new_tokens=20)
+        prompt_inputs = []
+        if isinstance(prompt, Prompt):
+            prompt_inputs = [prompt.prompt]
+        elif isinstance(prompt, list):
+            prompt_inputs = [p.prompt for p in prompt]
+            
+        logger.info(f"DefaultPipeline.streamGenerate with prompt_inputs: {prompt_inputs}")
+        input_ids = self.tokenizer(prompt_inputs, return_tensors="pt")
+        # input_ids = self.tokenizer([prompt], return_tensors="pt")
         max_new_tokens = 256
         if generate_kwargs["max_new_tokens"]:
             max_new_tokens = generate_kwargs["max_new_tokens"]

llmserve/backend/llm/pipelines/llamacpp/llamacpp_pipeline.py

@@ -4,7 +4,7 @@
 import torch
 
 from llmserve.backend.logger import get_logger
-from llmserve.backend.server.models import Response
+from llmserve.backend.server.models import Prompt, Response
 
 from ...initializers.llamacpp import LlamaCppInitializer, LlamaCppTokenizer
 from .._base import StreamingPipeline
@@ -225,7 +225,7 @@ def from_initializer(
             **kwargs,
         )
 
-    def streamGenerate(self, prompt: str, **generate_kwargs) -> Generator[str, None, None]:
+    def streamGenerate(self, prompt: Union[Prompt, List[Prompt]], **generate_kwargs) -> Generator[str, None, None]:
         logger.info(f"stream prompt: {prompt}")
         inputs = construct_prompts(prompt, prompt_format=self.prompt_format)
         logger.info(f"stream inputs: {inputs}")

llmserve/backend/llm/predictor.py

@@ -14,7 +14,7 @@
 from llmserve.backend.logger import get_logger
 from llmserve.backend.server.models import Args, Prompt
 
-from typing import AsyncGenerator, Generator
+from typing import AsyncGenerator, Generator, Union
 
 initialize_node_remote = ray.remote(initialize_node)
 logger = get_logger(__name__)
@@ -181,7 +181,7 @@ async def _predict_async(
     async def check_health(self):
         self.engine.check_health()
 
-    async def stream_generate_texts(self, prompt: str) -> Generator[str, None, None]: # type: ignore
+    async def stream_generate_texts(self, prompt: Union[Prompt, List[Prompt]]) -> Generator[str, None, None]: # type: ignore
         logger.info(f"call LLMPredictor.stream_generate_texts")
         for s in self.engine.stream_generate_texts(prompt):
             logger.info(f"LLMPredictor.stream_generate_texts -> yield ->{s}")

llmserve/backend/server/app.py

@@ -368,7 +368,7 @@ def streamer(self, model: str, prompt: Union[Prompt, List[Prompt]]) -> Streaming
         logger.info(f"search stream model key: {modelID}")
         return StreamingResponse(self.streamer_generate_text(modelID, prompt), media_type="text/plain")
 
-    async def streamer_generate_text(self, modelID: str, prompt: str) -> AsyncGenerator[str, None]:
+    async def streamer_generate_text(self, modelID: str, prompt: Union[Prompt, List[Prompt]]) -> AsyncGenerator[str, None]:
         logger.info(f'streamer_generate_text: {modelID}, prompt: "{prompt}"')
         r: DeploymentResponseGenerator = self._models[modelID].stream_generate_texts.remote(prompt)
         async for i in r: