Repository: hpcaitech/ColossalAI
File: colossalai/inference/config.py (Python)

Code before:
    def _verify_config(self) -> None:
        """
        Verify the input config
        """
        # check dtype
        if isinstance(self.dtype, str):
            # convert string dtype to torch dtype
            assert (
                self.dtype in _DTYPE_MAPPING
            ), f"Expected the dtype string argument to be in {list(_DTYPE_MAPPING.keys())} but found an unknown dtype: {self.dtype}"
            self.dtype = _DTYPE_MAPPING[self.dtype]
        assert (
            self.dtype in _ALLOWED_DTYPES
        ), f"Expected dtype to be in {_ALLOWED_DTYPES} but found an unknown dtype: {self.dtype}"

        # check distributed
        assert (
            self.tp_size * self.pp_size == dist.get_world_size()
        ), f"TP size({self.tp_size}) * PP size({self.pp_size}) should be equal to the global world size ({dist.get_world_size()})"

    def _get_tokenizer(self) -> None:
        if self.tokenizer is not None:
            return
        model_name = self.model.lower()
        if "llama" in model_name:
            self.tokenizer = "hf-internal-testing/llama-tokenizer"

    def _to_generation_config(self, model_config) -> GenerationConfig:
        meta_config = {
            "max_length": self.max_input_len + self.max_output_len,
            "max_new_tokens": self.max_output_len,
            "early_stopping": self.early_stopping,
            "do_sample": self.do_sample,
            "num_beams": self.beam_width,
        }
        for type in ["top_k", "top_p", "min_p"]:
            if hasattr(self, type):
                meta_config[type] = getattr(self, type)
        for type in ["pad_token_id", "bos_token_id", "eos_token_id"]:
            if hasattr(self, type):
                meta_config[type] = getattr(self, type)

        return GenerationConfig.from_dict(meta_config)


Review comment (suggestion):
You don't need to deal with tokenizer in `InferenceConfig`, always obey the `single-resposibility` principle. As I can see, the engine will receive a tokenizer, let the engine alone handle that, inference config should not handle tokenizer anymore to reduce coupling.

Code after:
    top_p: Optional[float] = None
    min_p: Optional[float] = None
    prompt_template: Optional[str] = None

    def __post_init__(self):
        self._verify_config()

    def _verify_config(self) -> None:
        """
        Verify the input config
        """
        # check dtype
        if isinstance(self.dtype, str):
            # convert string dtype to torch dtype
            assert (
                self.dtype in _DTYPE_MAPPING
            ), f"Expected the dtype string argument to be in {list(_DTYPE_MAPPING.keys())} but found an unknown dtype: {self.dtype}"
            self.dtype = _DTYPE_MAPPING[self.dtype]
        assert (
            self.dtype in _ALLOWED_DTYPES
        ), f"Expected dtype to be in {_ALLOWED_DTYPES} but found an unknown dtype: {self.dtype}"

        # check distributed
        assert (
            self.tp_size * self.pp_size == dist.get_world_size()
        ), f"TP size({self.tp_size}) * PP size({self.pp_size}) should be equal to the global world size ({dist.get_world_size()})"
        # check prompt template
        if self.prompt_template is None:
            return

        if self.prompt_template in _DEFAULT_PROMPT_TEMPLATES:
            self.prompt_template = _DEFAULT_PROMPT_TEMPLATES[self.prompt_template]
        else:
            # make sure the template can be formatted with input_text
            assert (
                "{input_text}" in self.prompt_template
            ), "The prompt template should contain '{input_text}' for formatting the input text. For example: 'USER: {input_text}\n\nASSISTANT: '"

    def to_generation_config(self, model_config) -> GenerationConfig:
        meta_config = {
            "max_length": self.max_input_len + self.max_output_len,
            "max_new_tokens": self.max_output_len,
            "early_stopping": self.early_stopping,
            "do_sample": self.do_sample,
            "num_beams": self.beam_width,
        }
        for type in ["top_k", "top_p", "min_p"]:
            if hasattr(self, type):
                meta_config[type] = getattr(self, type)
        for type in ["pad_token_id", "bos_token_id", "eos_token_id"]:
            if hasattr(model_config, type):