llm/Lib/site-packages/ctransformers/hub.py

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from huggingface_hub import HfApi, snapshot_download
from huggingface_hub.utils import validate_repo_id, HFValidationError

from .llm import Config, LLM


def get_path_type(path: str) -> Optional[str]:
    p = Path(path)
    if p.is_file():
        return "file"
    elif p.is_dir():
        return "dir"
    try:
        validate_repo_id(path)
        return "repo"
    except HFValidationError:
        pass


@dataclass
class AutoConfig:
    config: Config
    model_type: Optional[str] = None

    @classmethod
    def from_pretrained(
        cls,
        model_path_or_repo_id: str,
        local_files_only: bool = False,
        revision: Optional[str] = None,
        **kwargs,
    ) -> "AutoConfig":
        path_type = get_path_type(model_path_or_repo_id)
        if not path_type:
            raise ValueError(f"Model path '{model_path_or_repo_id}' doesn't exist.")

        config = Config()
        auto_config = AutoConfig(config=config)

        if path_type == "dir":
            cls._update_from_dir(model_path_or_repo_id, auto_config)
        elif path_type == "repo":
            cls._update_from_repo(
                model_path_or_repo_id,
                auto_config,
                local_files_only=local_files_only,
                revision=revision,
            )

        for k, v in kwargs.items():
            if not hasattr(config, k):
                raise TypeError(
                    f"'{k}' is an invalid keyword argument for from_pretrained()"
                )
            setattr(config, k, v)

        return auto_config

    @classmethod
    def _update_from_repo(
        cls,
        repo_id: str,
        auto_config: "AutoConfig",
        local_files_only: bool,
        revision: Optional[str] = None,
    ) -> None:
        path = snapshot_download(
            repo_id=repo_id,
            allow_patterns="config.json",
            local_files_only=local_files_only,
            revision=revision,
        )
        cls._update_from_dir(path, auto_config)

    @classmethod
    def _update_from_dir(cls, path: str, auto_config: "AutoConfig") -> None:
        path = (Path(path) / "config.json").resolve()
        if path.is_file():
            cls._update_from_file(path, auto_config)

    @classmethod
    def _update_from_file(cls, path: str, auto_config: "AutoConfig") -> None:
        with open(path) as f:
            config = json.load(f)

        auto_config.model_type = config.get("model_type")
        params = config.get("task_specific_params", {})
        params = params.get("text-generation", {})
        for name in [
            "top_k",
            "top_p",
            "temperature",
            "repetition_penalty",
            "last_n_tokens",
        ]:
            value = params.get(name)
            if value is not None:
                setattr(auto_config.config, name, value)


class AutoModelForCausalLM:
    @classmethod
    def from_pretrained(
        cls,
        model_path_or_repo_id: str,
        *,
        model_type: Optional[str] = None,
        model_file: Optional[str] = None,
        config: Optional[AutoConfig] = None,
        lib: Optional[str] = None,
        local_files_only: bool = False,
        revision: Optional[str] = None,
        hf: bool = False,
        **kwargs,
    ) -> LLM:
        """Loads the language model from a local file or remote repo.

        Args:
            model_path_or_repo_id: The path to a model file or directory or the
            name of a Hugging Face Hub model repo.
            model_type: The model type.
            model_file: The name of the model file in repo or directory.
            config: `AutoConfig` object.
            lib: The path to a shared library or one of `avx2`, `avx`, `basic`.
            local_files_only: Whether or not to only look at local files
            (i.e., do not try to download the model).
            revision: The specific model version to use. It can be a branch
            name, a tag name, or a commit id.
            hf: Whether to create a Hugging Face Transformers model.

        Returns:
            `LLM` object.
        """
        if model_type is None and "gptq" in str(model_path_or_repo_id).lower():
            model_type = "gptq"
        if model_type == "gptq":
            from . import gptq

            return gptq.AutoModelForCausalLM.from_pretrained(
                model_path_or_repo_id,
                local_files_only=local_files_only,
                revision=revision,
                **kwargs,
            )

        config = config or AutoConfig.from_pretrained(
            model_path_or_repo_id,
            local_files_only=local_files_only,
            revision=revision,
            **kwargs,
        )
        model_type = model_type or config.model_type

        path_type = get_path_type(model_path_or_repo_id)
        model_path = None
        if path_type == "file":
            model_path = model_path_or_repo_id
        elif path_type == "dir":
            model_path = cls._find_model_path_from_dir(
                model_path_or_repo_id, model_file
            )
        elif path_type == "repo":
            model_path = cls._find_model_path_from_repo(
                model_path_or_repo_id,
                model_file,
                local_files_only=local_files_only,
                revision=revision,
            )

        llm = LLM(
            model_path=model_path,
            model_type=model_type,
            config=config.config,
            lib=lib,
        )
        if not hf:
            return llm

        from .transformers import CTransformersConfig, CTransformersModel

        config = CTransformersConfig(name_or_path=str(model_path_or_repo_id))
        return CTransformersModel(config=config, llm=llm)

    @classmethod
    def _find_model_path_from_repo(
        cls,
        repo_id: str,
        filename: Optional[str],
        local_files_only: bool,
        revision: Optional[str] = None,
    ) -> str:
        if not filename and not local_files_only:
            filename = cls._find_model_file_from_repo(
                repo_id=repo_id,
                revision=revision,
            )
        allow_patterns = filename or ["*.bin", "*.gguf"]
        path = snapshot_download(
            repo_id=repo_id,
            allow_patterns=allow_patterns,
            local_files_only=local_files_only,
            revision=revision,
        )
        return cls._find_model_path_from_dir(path, filename=filename)

    @classmethod
    def _find_model_file_from_repo(
        cls,
        repo_id: str,
        revision: Optional[str] = None,
    ) -> Optional[str]:
        api = HfApi()
        repo_info = api.repo_info(
            repo_id=repo_id,
            files_metadata=True,
            revision=revision,
        )
        files = [
            (f.size, f.rfilename)
            for f in repo_info.siblings
            if f.rfilename.endswith(".bin") or f.rfilename.endswith(".gguf")
        ]
        if not files:
            raise ValueError(f"No model file found in repo '{repo_id}'")
        return min(files)[1]

    @classmethod
    def _find_model_path_from_dir(
        cls,
        path: str,
        filename: Optional[str] = None,
    ) -> str:
        path = Path(path).resolve()
        if filename:
            file = (path / filename).resolve()
            if not file.is_file():
                raise ValueError(f"Model file '{filename}' not found in '{path}'")
            return str(file)

        files = [
            (f.stat().st_size, f)
            for f in path.iterdir()
            if f.is_file() and (f.name.endswith(".bin") or f.name.endswith(".gguf"))
        ]
        if not files:
            raise ValueError(f"No model file found in directory '{path}'")
        file = min(files)[1]
        return str(file.resolve())


class AutoTokenizer:
    @classmethod
    def from_pretrained(cls, model):
        from .transformers import CTransformersModel, CTransformersTokenizer

        if not isinstance(model, CTransformersModel):
            raise TypeError(
                f"Currently `AutoTokenizer.from_pretrained` only accepts a model object. Please use:\n\n"
                "  model = AutoModelForCausalLM.from_pretrained(..., hf=True)\n"
                "  tokenizer = AutoTokenizer.from_pretrained(model)"
            )

        return CTransformersTokenizer(model._llm)