lmdeploy_engine.py

# Copyright (c) ModelScope Contributors. All rights reserved.
import asyncio
import inspect
import lmdeploy
import os
import time
import torch
from contextlib import contextmanager
from copy import deepcopy
from lmdeploy import PytorchEngineConfig, TurbomindEngineConfig, VisionConfig, pipeline
from lmdeploy.api import autoget_backend_config
from lmdeploy.serve import async_engine
from packaging import version
from PIL import Image
from transformers import GenerationConfig
from transformers.utils.versions import require_version
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union

from swift.metrics import Metric
from swift.model import get_processor
from swift.template import Template
from swift.utils import get_logger, get_seed, safe_snapshot_download
from .infer_engine import InferEngine
from .patch import patch_auto_config, patch_auto_tokenizer
from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
                       ChatCompletionStreamResponse, ChatMessage, DeltaMessage, InferRequest, RequestConfig)
from .utils import InferStreamer

try:
    from lmdeploy import EngineGenerationConfig as LmdeployGenerationConfig
except ImportError:
    # compat lmdeploy >= 0.6.*
    from lmdeploy import GenerationConfig as LmdeployGenerationConfig

logger = get_logger()


class LmdeployEngine(InferEngine):

    def __init__(
        self,
        model_id_or_path: str,
        *,
        template: Optional[Template] = None,
        torch_dtype: Optional[torch.dtype] = None,
        model_type: Optional[str] = None,
        template_type: Optional[str] = None,
        use_hf: Optional[bool] = None,
        hub_token: Optional[str] = None,
        revision: Optional[str] = None,
        # engine_kwargs
        tp: int = 1,
        session_len: Optional[int] = None,
        cache_max_entry_count: float = 0.8,
        quant_policy: int = 0,  # e.g. 4, 8
        vision_batch_size: int = 1,  # max_batch_size in VisionConfig
        engine_kwargs: Optional[Dict[str, Any]] = None,
        devices: Optional[List[int]] = None,
    ) -> None:
        self.model_id_or_path = model_id_or_path
        self.torch_dtype = torch_dtype
        self.model_type = model_type
        self.use_hf = use_hf
        self.hub_token = hub_token
        self.revision = revision
        self.tp = tp
        self.session_len = session_len
        self.cache_max_entry_count = cache_max_entry_count
        self.quant_policy = quant_policy
        self.vision_batch_size = vision_batch_size
        self.devices = devices
        if template is None:
            processor = self._get_processor()
            template = self._get_template(processor, template_type=template_type)
        else:
            safe_snapshot_download(
                model_id_or_path,
                revision=revision,
                download_model=True,
                use_hf=use_hf,
                ignore_patterns=getattr(template.model_meta, 'ignore_patterns', None),
                hub_token=hub_token)
        super().__init__(template)

        if self.max_model_len is not None:
            self.max_model_len -= 1
        self._prepare_engine_kwargs(engine_kwargs)
        self.config.torch_dtype = self.torch_dtype = self.torch_dtype or self.model_info.torch_dtype
        self._prepare_engine()
        self._load_generation_config()

    def _get_processor(self):
        return get_processor(
            model_id_or_path=self.model_id_or_path,
            torch_dtype=self.torch_dtype,
            download_model=True,
            model_type=self.model_type,
            use_hf=self.use_hf,
            hub_token=self.hub_token,
            revision=self.revision)

    def _prepare_engine_kwargs(self, engine_kwargs):
        if engine_kwargs is None:
            engine_kwargs = {}
        engine_kwargs['tp'] = self.tp
        engine_kwargs['session_len'] = self.session_len
        engine_kwargs['cache_max_entry_count'] = self.cache_max_entry_count
        engine_kwargs['quant_policy'] = self.quant_policy
        if 'devices' in inspect.signature(TurbomindEngineConfig).parameters:
            engine_kwargs['devices'] = self.devices
        backend_config = TurbomindEngineConfig(**engine_kwargs)
        backend_config = autoget_backend_config(self.model_dir, backend_config)
        self.backend_config = backend_config
        logger.info(f'backend_config: {backend_config}')

        pipeline_kwargs = {}
        is_multimodal = self.model_meta.is_multimodal
        if is_multimodal:
            require_version(
                'lmdeploy<0.9', 'LmdeployEngine will no longer maintain inference for '
                'multimodal models in lmdeploy>=0.9.')
            vision_config = VisionConfig(max_batch_size=self.vision_batch_size)
            pipeline_kwargs['vision_config'] = vision_config
            logger.info(f'vision_config: {vision_config}')
        self.pipeline_kwargs = pipeline_kwargs

    @contextmanager
    def _patch_pipeline(self):
        _old_best_match_model = async_engine.best_match_model

        def _best_match_model(*args, **kwargs) -> Optional[str]:
            return self.model_info.model_type

        async_engine.best_match_model = _best_match_model
        try:
            yield
        finally:
            async_engine.best_match_model = _old_best_match_model

    def _prepare_engine(self):
        with patch_auto_tokenizer(self.tokenizer), patch_auto_config(self.config), self._patch_pipeline():
            engine = pipeline(self.model_dir, backend_config=self.backend_config, **self.pipeline_kwargs)
        self.engine = engine

    def _load_generation_config(self):
        generation_config_path = os.path.join(self.model_dir, 'generation_config.json')
        if os.path.isfile(generation_config_path):
            generation_config = GenerationConfig.from_pretrained(self.model_dir)
            kwargs = generation_config.to_dict()
            max_new_tokens = kwargs.get('max_new_tokens')
            if max_new_tokens is None:
                kwargs.pop('max_new_tokens', None)
            parameters = inspect.signature(LmdeployGenerationConfig).parameters
            for k, v in kwargs.copy().items():
                if k not in parameters or v is None:
                    kwargs.pop(k)
            self.generation_config = LmdeployGenerationConfig(**kwargs)
        else:
            self.generation_config = LmdeployGenerationConfig()

    def _add_stop_words(self, generation_config: LmdeployGenerationConfig, request_config: RequestConfig) -> None:
        template_meta = self.template.template_meta
        stop_words = (request_config.stop or []) + (self.generation_config.stop_words or []) + template_meta.stop_words
        generation_config.stop_words = self._get_stop_token_ids(stop_words)
        # compat lmdeploy >= 0.6.*
        generation_config.stop_token_ids = generation_config.stop_words

    def _prepare_generation_config(self, request_config: RequestConfig) -> LmdeployGenerationConfig:
        kwargs = {'max_new_tokens': request_config.max_tokens}
        for key in ['temperature', 'top_k', 'top_p', 'repetition_penalty']:
            new_value = getattr(request_config, key)
            if new_value is None:
                kwargs[key] = getattr(self.generation_config, key)
            else:
                kwargs[key] = new_value
        if request_config.seed is None:
            request_config.seed = get_seed()
        kwargs['random_seed'] = request_config.seed
        if request_config.temperature == 0:
            kwargs['temperature'] = 1  # avoid unnecessary process
            kwargs['top_k'] = 1

        if request_config.logprobs:
            kwargs['logprobs'] = 1
            if request_config.top_logprobs is not None:
                kwargs['logprobs'] = max(1, request_config.top_logprobs)

        res = LmdeployGenerationConfig(**kwargs)
        return res

    async def _infer_stream_async(
        self,
        inputs: Dict[str, Any],
        generation_config: LmdeployGenerationConfig,
        request_config: RequestConfig,
    ) -> AsyncIterator[ChatCompletionStreamResponse]:
        session_id = time.time_ns()
        kwargs = {'stream_output': True, 'gen_config': generation_config, 'sequence_start': True, 'sequence_end': True}
        if version.parse(lmdeploy.__version__) >= version.parse('0.6.5'):
            async with self.engine.model_inst(session_id) as inst:
                context = self.engine.safe_run(inst, session_id, **inputs, **kwargs)
        else:
            context = self.engine.safe_run(session_id)

        infer_streamer = InferStreamer(self.template, template_inputs=inputs['template_inputs'])
        token_idx = 0
        async with context as gen:
            if version.parse(lmdeploy.__version__) < version.parse('0.6.5'):
                generator = await self.engine.get_generator(False, session_id)
                gen = generator.async_stream_infer(session_id=session_id, **inputs, **kwargs)
            is_finished = False
            while not is_finished:
                try:
                    output = await gen.__anext__()
                except StopAsyncIteration:
                    is_finished = True
                delta_text = infer_streamer.get_printable_text(output.token_ids, is_finished)
                if not delta_text and not is_finished:
                    continue

                logprobs = self._get_logprobs(output.logprobs, output.token_ids[token_idx:],
                                              request_config.top_logprobs)
                token_idx = len(output.token_ids)

                usage_info = self._get_usage_info(len(inputs['input_ids']), output.num_token)
                toolcall = None
                if is_finished:
                    toolcall = self._get_toolcall(
                        self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs']))
                finish_reason = self._get_finish_reason(generation_config.max_new_tokens, output.num_token,
                                                        output.status.name == 'FINISH')
                choices = [
                    ChatCompletionResponseStreamChoice(
                        index=0,
                        delta=DeltaMessage(role='assistant', content=delta_text, tool_calls=toolcall),
                        finish_reason=finish_reason,
                        logprobs=logprobs)
                ]
                yield ChatCompletionStreamResponse(model=self.model_name, choices=choices, usage=usage_info)

    async def _infer_full_async(
        self,
        inputs: Dict[str, Any],
        generation_config: LmdeployGenerationConfig,
        request_config: RequestConfig,
    ) -> ChatCompletionResponse:
        session_id = time.time_ns()
        kwargs = {'stream_output': False, 'gen_config': generation_config, 'sequence_start': True, 'sequence_end': True}
        if version.parse(lmdeploy.__version__) >= version.parse('0.6.5'):
            async with self.engine.model_inst(session_id) as inst:
                async with self.engine.safe_run(inst, session_id, **inputs, **kwargs) as gen:
                    async for output in gen:
                        pass
                if self.engine.backend == 'pytorch':
                    # manually end pytorch session
                    await inst.async_end(session_id)

        else:
            async with self.engine.safe_run(session_id):
                generator = await self.engine.get_generator(False, session_id)
                async for output in generator.async_stream_infer(session_id=session_id, **inputs, **kwargs):
                    pass

        response = self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs'])
        logprobs = self._get_logprobs(output.logprobs, output.token_ids, request_config.top_logprobs)

        usage_info = self._get_usage_info(len(inputs['input_ids']), output.num_token)
        toolcall = self._get_toolcall(response)
        finish_reason = self._get_finish_reason(generation_config.max_new_tokens, output.num_token,
                                                output.status.name == 'FINISH')
        token_ids = output.token_ids if request_config.return_details else None
        choices = [
            ChatCompletionResponseChoice(
                index=0,
                message=ChatMessage(role='assistant', content=response, tool_calls=toolcall),
                finish_reason=finish_reason,
                logprobs=logprobs,
                token_ids=token_ids)
        ]
        prompt_token_ids = None
        images_size = None
        if request_config.return_details:
            prompt_token_ids = inputs['input_ids']
            images = inputs['template_inputs'].images
            if all(isinstance(image, Image.Image) for image in images):
                images_size = [image.size for image in images]
        return ChatCompletionResponse(
            model=self.model_name,
            choices=choices,
            usage=usage_info,
            prompt_token_ids=prompt_token_ids,
            images_size=images_size)

    async def infer_async(self,
                          infer_request: InferRequest,
                          request_config: Optional[RequestConfig] = None,
                          *,
                          pre_infer_hook=None,
                          **kwargs) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]:
        request_config = deepcopy(request_config or RequestConfig())
        self.template.set_mode('lmdeploy')

        loop = asyncio.get_running_loop()
        with torch.inference_mode():
            inputs = await loop.run_in_executor(None, self.template.encode, infer_request, True)
        images = inputs.pop('images', None)
        if images:
            if version.parse(lmdeploy.__version__) >= version.parse('0.6.5'):
                messages = self.engine._convert_prompts(('', images))
                messages = await self.engine.async_convert_to_pil_images(messages)
                results = await self.engine.vl_encoder.preprocess(messages)
                if self.engine.backend == 'turbomind':
                    results = await self.engine.vl_encoder.async_infer(results)
                    inputs['images'] = [result['content'] for result in results if result['role'] == 'forward'][0]
                    await self.template.prepare_lmdeploy_turbomind_inputs(inputs)
                else:
                    inputs['images'] = results[1]['content']
                    await self.template.prepare_lmdeploy_pytorch_inputs(inputs)
            else:
                inputs['images'] = await self.engine.vl_encoder.async_infer(images)
                await self.template.prepare_lmdeploy_turbomind_inputs(inputs)

        self.set_default_max_tokens(request_config, inputs)
        generation_config = self._prepare_generation_config(request_config)
        self._add_stop_words(generation_config, request_config)
        kwargs.update({'inputs': inputs, 'generation_config': generation_config, 'request_config': request_config})
        if pre_infer_hook:
            kwargs = pre_infer_hook(kwargs)
        if request_config.stream:
            return self._infer_stream_async(**kwargs)
        else:
            return await self._infer_full_async(**kwargs)

    def _batch_infer_stream(self, *args, **kwargs):
        if hasattr(self.engine, 'vl_encoder'):
            self.engine.vl_encoder._loop_task = None
        if hasattr(self.engine, 'free_insts'):
            self.engine.free_insts = None
        return super()._batch_infer_stream(*args, **kwargs)

    def infer(
        self,
        infer_requests: List[InferRequest],
        request_config: Optional[RequestConfig] = None,
        metrics: Optional[List[Metric]] = None,
        *,
        use_tqdm: Optional[bool] = None,
        **kwargs,
    ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
        return super().infer(infer_requests, request_config, metrics, use_tqdm=use_tqdm, **kwargs)