codekingpro/portable-devtools

1import logging2from typing import Any, Optional3 4from langchain_core.language_models.llms import LLM5 6from langchain_community.llms.ipex_llm import IpexLLM7 8logger = logging.getLogger(__name__)9 10 11class BigdlLLM(IpexLLM):12    """Wrapper around the BigdlLLM model13 14    Example:15        .. code-block:: python16 17            from langchain_community.llms import BigdlLLM18            llm = BigdlLLM.from_model_id(model_id="THUDM/chatglm-6b")19    """20 21    @classmethod22    def from_model_id(23        cls,24        model_id: str,25        model_kwargs: Optional[dict] = None,26        *,27        tokenizer_id: Optional[str] = None,28        load_in_4bit: bool = True,29        load_in_low_bit: Optional[str] = None,30        **kwargs: Any,31    ) -> LLM:32        """33        Construct object from model_id34 35        Args:36            model_id: Path for the huggingface repo id to be downloaded or37                      the huggingface checkpoint folder.38            tokenizer_id: Path for the huggingface repo id to be downloaded or39                      the huggingface checkpoint folder which contains the tokenizer.40            model_kwargs: Keyword arguments to pass to the model and tokenizer.41            kwargs: Extra arguments to pass to the model and tokenizer.42 43        Returns:44            An object of BigdlLLM.45        """46        logger.warning("BigdlLLM was deprecated. Please use IpexLLM instead.")47 48        try:49            from bigdl.llm.transformers import (50                AutoModel,51                AutoModelForCausalLM,52            )53            from transformers import AutoTokenizer, LlamaTokenizer54 55        except ImportError:56            raise ImportError(57                "Could not import bigdl-llm or transformers. "58                "Please install it with `pip install --pre --upgrade bigdl-llm[all]`."59            )60 61        if load_in_low_bit is not None:62            logger.warning(63                """`load_in_low_bit` option is not supported in BigdlLLM and 64                is ignored. For more data types support with `load_in_low_bit`, 65                use IpexLLM instead."""66            )67 68        if not load_in_4bit:69            raise ValueError(70                "BigdlLLM only supports loading in 4-bit mode, "71                "i.e. load_in_4bit = True. "72                "Please install it with `pip install --pre --upgrade bigdl-llm[all]`."73            )74 75        _model_kwargs = model_kwargs or {}76        _tokenizer_id = tokenizer_id or model_id77 78        try:79            tokenizer = AutoTokenizer.from_pretrained(_tokenizer_id, **_model_kwargs)80        except Exception:81            tokenizer = LlamaTokenizer.from_pretrained(_tokenizer_id, **_model_kwargs)82 83        try:84            model = AutoModelForCausalLM.from_pretrained(85                model_id, load_in_4bit=True, **_model_kwargs86            )87        except Exception:88            model = AutoModel.from_pretrained(89                model_id, load_in_4bit=True, **_model_kwargs90            )91 92        if "trust_remote_code" in _model_kwargs:93            _model_kwargs = {94                k: v for k, v in _model_kwargs.items() if k != "trust_remote_code"95            }96 97        return cls(98            model_id=model_id,99            model=model,100            tokenizer=tokenizer,101            model_kwargs=_model_kwargs,102            **kwargs,103        )104 105    @classmethod106    def from_model_id_low_bit(107        cls,108        model_id: str,109        model_kwargs: Optional[dict] = None,110        *,111        tokenizer_id: Optional[str] = None,112        **kwargs: Any,113    ) -> LLM:114        """115        Construct low_bit object from model_id116 117        Args:118 119            model_id: Path for the bigdl-llm transformers low-bit model folder.120            tokenizer_id: Path for the huggingface repo id or local model folder121                      which contains the tokenizer.122            model_kwargs: Keyword arguments to pass to the model and tokenizer.123            kwargs: Extra arguments to pass to the model and tokenizer.124 125        Returns:126            An object of BigdlLLM.127        """128 129        logger.warning("BigdlLLM was deprecated. Please use IpexLLM instead.")130 131        try:132            from bigdl.llm.transformers import (133                AutoModel,134                AutoModelForCausalLM,135            )136            from transformers import AutoTokenizer, LlamaTokenizer137 138        except ImportError:139            raise ImportError(140                "Could not import bigdl-llm or transformers. "141                "Please install it with `pip install --pre --upgrade bigdl-llm[all]`."142            )143 144        _model_kwargs = model_kwargs or {}145        _tokenizer_id = tokenizer_id or model_id146 147        try:148            tokenizer = AutoTokenizer.from_pretrained(_tokenizer_id, **_model_kwargs)149        except Exception:150            tokenizer = LlamaTokenizer.from_pretrained(_tokenizer_id, **_model_kwargs)151 152        try:153            model = AutoModelForCausalLM.load_low_bit(model_id, **_model_kwargs)154        except Exception:155            model = AutoModel.load_low_bit(model_id, **_model_kwargs)156 157        if "trust_remote_code" in _model_kwargs:158            _model_kwargs = {159                k: v for k, v in _model_kwargs.items() if k != "trust_remote_code"160            }161 162        return cls(163            model_id=model_id,164            model=model,165            tokenizer=tokenizer,166            model_kwargs=_model_kwargs,167            **kwargs,168        )169 170    @property171    def _llm_type(self) -> str:172        return "bigdl-llm"173