""" # Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. # # 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. """ from __future__ import annotations import sys import traceback import uuid import time from typing import Optional, Dict, List, Any, Union, overload from tqdm import tqdm from fastdeploy.engine.args_utils import EngineArgs from fastdeploy.engine.engine import LLMEngine from fastdeploy.engine.sampling_params import SamplingParams from fastdeploy.entrypoints.chat_utils import ChatCompletionMessageParam from fastdeploy.utils import llm_logger import logging root_logger = logging.getLogger() for handler in root_logger.handlers[:]: if isinstance(handler, logging.StreamHandler): root_logger.removeHandler(handler) class LLM: """ Language Model wrapper class providing high-level interfaces for text generation. This class manages the LLMEngine instance and provides convenient methods for generating text and chat completions. Attributes: llm_engine: Underlying LLMEngine instance default_sampling_params: Default sampling parameters for generation Args: model: Name of the language model to use tokenizer: Name of the tokenizer to use (defaults to model's tokenizer) **kwargs: Additional arguments passed to EngineArgs constructor Raises: ValueError: If model is not supported RuntimeError: If engine fails to start """ def __init__( self, model: str, tokenizer: Optional[str] = None, **kwargs, ): engine_args = EngineArgs( model=model, tokenizer=tokenizer, **kwargs, ) # Create the Engine self.llm_engine = LLMEngine.from_engine_args( engine_args=engine_args) self.default_sampling_params = SamplingParams( max_tokens=self.llm_engine.cfg.max_model_len) self.llm_engine.start() def generate( self, prompts: Union[str, list[str], list[int], list[list[int]], dict[str, Any], list[dict[str, Any]]], sampling_params: Optional[Union[SamplingParams, list[SamplingParams]]] = None, use_tqdm: bool = True, ): """ Generate text based on input prompts. Supports various input formats including: - Single prompt string - List of prompt strings - Token IDs (single or batched) - Dictionary with additional parameters - List of parameter dictionaries Args: prompts: Input prompts in various formats sampling_params: Sampling parameters for generation use_tqdm: Whether to show progress bar Returns: Generated text output(s) Raises: ValueError: If prompts and sampling_params length mismatch TypeError: If prompts format is invalid """ if sampling_params is None: sampling_params = self.default_sampling_params if isinstance(sampling_params, SamplingParams): sampling_params_len = 1 else: sampling_params_len = len(sampling_params) if isinstance(prompts, str): prompts = [prompts] if isinstance(prompts, list) and isinstance(prompts[0], int): prompts = [prompts] if isinstance(prompts, dict): if "prompts" not in prompts: raise ValueError("prompts must be a input dict") prompts = [prompts] sampling_params = None if sampling_params_len != 1 and len(prompts) != sampling_params_len: raise ValueError( "prompts and sampling_params must be the same length.") req_ids = self._add_request( prompts=prompts, sampling_params=sampling_params ) # get output outputs = self._run_engine(req_ids, use_tqdm=use_tqdm) return outputs def chat( self, messages: Union[list[ChatCompletionMessageParam], list[list[ChatCompletionMessageParam]]], sampling_params: Optional[Union[SamplingParams, list[SamplingParams]]] = None, use_tqdm: bool = True, ): """ Generate chat completions based on conversation messages. Args: messages: Single conversation or list of conversations sampling_params: Sampling parameters for generation use_tqdm: Whether to show progress bar Returns: Generated chat response(s) Raises: ValueError: If messages and sampling_params length mismatch """ if sampling_params is None: sampling_params = self.default_sampling_params if isinstance(sampling_params, SamplingParams): sampling_params_len = 1 else: sampling_params_len = len(sampling_params) if isinstance(messages, list) and isinstance(messages[0], dict): messages = [messages] if sampling_params_len != 1 and len(messages) != sampling_params_len: raise ValueError( "messages and sampling_params must be the same length.") messages_len = len(messages) for i in range(messages_len): messages[i] = { "messages": messages[i] } req_ids = self._add_request( prompts=messages, sampling_params=sampling_params ) # get output outputs = self._run_engine(req_ids, use_tqdm=use_tqdm) return outputs def _add_request( self, prompts, sampling_params, ): """ Add generation requests to the LLM engine. Args: prompts: Input prompts to process sampling_params: Sampling parameters for generation Returns: list: List of generated request IDs Raises: ValueError: If prompts is None TypeError: If prompts format is invalid """ if prompts is None: raise ValueError("prompts and prompt_ids cannot be both None.") prompts_len = len(prompts) req_ids = [] for i in range(prompts_len): request_id = str(uuid.uuid4()) if isinstance(prompts[i], str): tasks = { "prompt": prompts[i], "request_id": request_id, } elif isinstance(prompts[i], list) and isinstance(prompts[i][0], int): tasks = { "prompt_token_ids": prompts[i], "request_id": request_id, } elif isinstance(prompts[i], dict): tasks = prompts[i] tasks["request_id"] = request_id else: raise TypeError( f"Invalid type for 'prompt': {type(prompts[i])}, expected one of ['str', 'list', 'dict']." ) req_ids.append(request_id) if isinstance(sampling_params, list): sampling_params = sampling_params[i] self.llm_engine.add_requests(tasks, sampling_params) return req_ids def _run_engine( self, req_ids: list[str], use_tqdm: bool ): """ Run the engine and collect results for given request IDs. Args: req_ids: List of request IDs to process use_tqdm: Whether to show progress bar Returns: list: List of generation results Note: This method blocks until all requests are completed """ # Initialize tqdm. if use_tqdm: num_requests = len(req_ids) pbar = tqdm( total=num_requests, desc="Processed prompts", dynamic_ncols=True, postfix=(f"est. speed input: {0:.2f} toks/s, " f"output: {0:.2f} toks/s"), ) output = [] while num_requests: finished = [] for i, req_id in enumerate(req_ids): try: for result in self.llm_engine._get_generated_result(req_id): result = self.llm_engine.data_processor.process_response( result) llm_logger.debug( f"Send result to client under push mode: {result}") if result.finished: output.append(result) finished.append(i) llm_logger.debug( "Request id: {} has been completed.".format(req_id)) if use_tqdm: pbar.update(1) except Exception as e: llm_logger.error("Unexcepted error happend: {}".format(e)) num_requests -= len(finished) for i in reversed(finished): req_ids.pop(i) if use_tqdm: pbar.close() return output if __name__ == "__main__": # Example usage: # llm = LLM(model="llama_model") # output = llm.generate(prompts="who are you?", use_tqdm=True) # print(output) llm = LLM(model="/opt/baidu/paddle_internal/FastDeploy/Qwen2.5-7B", tensor_parallel_size=2) sampling_params = SamplingParams(temperature=0.1, max_tokens=30) output = llm.generate(prompts="who are you？", use_tqdm=True, sampling_params=sampling_params) print(output) output = llm.generate(prompts=["who are you？", "what can you do？"], sampling_params=SamplingParams( temperature=1, max_tokens=50), use_tqdm=True) print(output) output = llm.generate(prompts=["who are you？", "I miss you"], sampling_params=[SamplingParams( temperature=1, max_tokens=50), SamplingParams(temperature=1, max_tokens=20)], use_tqdm=True) print(output)