Class InferenceEngine

colossalai/inference/core/engine.py:16–133

InferenceEngine which manages the inference process.. Args: model_or_path (nn.Module or DiffusionPipeline or str): Path or nn.Module or DiffusionPipeline of this model. tokenizer Optional[(Union[PreTrainedTokenizer, PreTrainedTokenizerFast])]: Path of the tokenizer to use.

class InferenceEngine:
    """
    InferenceEngine which manages the inference process..

    Args:
        model_or_path (nn.Module or DiffusionPipeline or str): Path or nn.Module or DiffusionPipeline of this model.
        tokenizer Optional[(Union[PreTrainedTokenizer, PreTrainedTokenizerFast])]: Path of the tokenizer to use.
        inference_config (Optional[InferenceConfig], optional): Store the configuration information related to inference.
        verbose (bool): Determine whether or not to log the generation process.
        model_policy ("Policy"): the policy to shardformer model. It will be determined by the model type if not provided.
    """

    def __init__(
        self,
        model_or_path: Union[nn.Module, str, DiffusionPipeline],
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast] = None,
        inference_config: InferenceConfig = None,
        verbose: bool = False,
        model_policy: Union[Policy, Type[Policy]] = None,
    ) -> None:
        self.__dict__["_initialized"] = False  # use __dict__ directly to avoid calling __setattr__
        self.model_type = get_model_type(model_or_path=model_or_path)
        self.engine = None
        if self.model_type == ModelType.LLM:
            from .llm_engine import LLMEngine

            self.engine = LLMEngine(
                model_or_path=model_or_path,
                tokenizer=tokenizer,
                inference_config=inference_config,
                verbose=verbose,
                model_policy=model_policy,
            )
        elif self.model_type == ModelType.DIFFUSION_MODEL:
            from .diffusion_engine import DiffusionEngine

            self.engine = DiffusionEngine(
                model_or_path=model_or_path,
                inference_config=inference_config,
                verbose=verbose,
                model_policy=model_policy,
            )
        elif self.model_type == ModelType.UNKNOWN:
            self.logger.error(f"Model Type either Difffusion or LLM!")

        self._initialized = True
        self._verify_args()

    def _verify_args(self) -> None:
        """Verify the input args"""
        assert self.engine is not None, "Please init Engine first"
        assert self._initialized, "Engine must be initialized"

    def generate(
        self,
        request_ids: Union[List[int], int] = None,
        prompts: Union[List[str], str] = None,
        *args,