Source code for langchain_community.llms.bigdl_llm

import logging
from typing import Any, Optional

from langchain_core.language_models.llms import LLM

from langchain_community.llms.ipex_llm import IpexLLM

logger = logging.getLogger(__name__)


class BigdlLLM(IpexLLM):
    """包装器，用于BigdlLLM模型

    示例:
        .. code-block:: python

            from langchain_community.llms import BigdlLLM
            llm = BigdlLLM.from_model_id(model_id="THUDM/chatglm-6b")"""

    @classmethod
    def from_model_id(
        cls,
        model_id: str,
        model_kwargs: Optional[dict] = None,
        *,
        tokenizer_id: Optional[str] = None,
        load_in_4bit: bool = True,
        load_in_low_bit: Optional[str] = None,
        **kwargs: Any,
    ) -> LLM:
        """根据model_id构建对象

参数:
    model_id: 下载huggingface repo id或huggingface checkpoint文件夹的路径。
    tokenizer_id: 下载huggingface repo id或包含分词器的huggingface checkpoint文件夹的路径。
    model_kwargs: 传递给模型和分词器的关键字参数。
    kwargs: 传递给模型和分词器的额外参数。

返回:
    一个BigdlLLM对象。
"""
        logger.warning("BigdlLLM was deprecated. Please use IpexLLM instead.")

        try:
            from bigdl.llm.transformers import (
                AutoModel,
                AutoModelForCausalLM,
            )
            from transformers import AutoTokenizer, LlamaTokenizer

        except ImportError:
            raise ImportError(
                "Could not import bigdl-llm or transformers. "
                "Please install it with `pip install --pre --upgrade bigdl-llm[all]`."
            )

        if load_in_low_bit is not None:
            logger.warning(
                """`load_in_low_bit` option is not supported in BigdlLLM and 
                is ignored. For more data types support with `load_in_low_bit`, 
                use IpexLLM instead."""
            )

        if not load_in_4bit:
            raise ValueError(
                "BigdlLLM only supports loading in 4-bit mode, "
                "i.e. load_in_4bit = True. "
                "Please install it with `pip install --pre --upgrade bigdl-llm[all]`."
            )

        _model_kwargs = model_kwargs or {}
        _tokenizer_id = tokenizer_id or model_id

        try:
            tokenizer = AutoTokenizer.from_pretrained(_tokenizer_id, **_model_kwargs)
        except Exception:
            tokenizer = LlamaTokenizer.from_pretrained(_tokenizer_id, **_model_kwargs)

        try:
            model = AutoModelForCausalLM.from_pretrained(
                model_id, load_in_4bit=True, **_model_kwargs
            )
        except Exception:
            model = AutoModel.from_pretrained(
                model_id, load_in_4bit=True, **_model_kwargs
            )

        if "trust_remote_code" in _model_kwargs:
            _model_kwargs = {
                k: v for k, v in _model_kwargs.items() if k != "trust_remote_code"
            }

        return cls(
            model_id=model_id,
            model=model,
            tokenizer=tokenizer,
            model_kwargs=_model_kwargs,
            **kwargs,
        )

    @classmethod
    def from_model_id_low_bit(
        cls,
        model_id: str,
        model_kwargs: Optional[dict] = None,
        *,
        tokenizer_id: Optional[str] = None,
        **kwargs: Any,
    ) -> LLM:
        """从model_id构建low_bit对象

参数：

    model_id：bigdl-llm transformers low-bit模型文件夹的路径。
    tokenizer_id：包含分词器的huggingface repo id或本地模型文件夹的路径。
    model_kwargs：传递给模型和分词器的关键字参数。
    kwargs：传递给模型和分词器的额外参数。

返回：
    一个BigdlLLM对象。
"""

        logger.warning("BigdlLLM was deprecated. Please use IpexLLM instead.")

        try:
            from bigdl.llm.transformers import (
                AutoModel,
                AutoModelForCausalLM,
            )
            from transformers import AutoTokenizer, LlamaTokenizer

        except ImportError:
            raise ImportError(
                "Could not import bigdl-llm or transformers. "
                "Please install it with `pip install --pre --upgrade bigdl-llm[all]`."
            )

        _model_kwargs = model_kwargs or {}
        _tokenizer_id = tokenizer_id or model_id

        try:
            tokenizer = AutoTokenizer.from_pretrained(_tokenizer_id, **_model_kwargs)
        except Exception:
            tokenizer = LlamaTokenizer.from_pretrained(_tokenizer_id, **_model_kwargs)

        try:
            model = AutoModelForCausalLM.load_low_bit(model_id, **_model_kwargs)
        except Exception:
            model = AutoModel.load_low_bit(model_id, **_model_kwargs)

        if "trust_remote_code" in _model_kwargs:
            _model_kwargs = {
                k: v for k, v in _model_kwargs.items() if k != "trust_remote_code"
            }

        return cls(
            model_id=model_id,
            model=model,
            tokenizer=tokenizer,
            model_kwargs=_model_kwargs,
            **kwargs,
        )

    @property
    def _llm_type(self) -> str:
        return "bigdl-llm"