Draken007's picture
Upload 7228 files
2a0bc63 verified
from pathlib import Path
from typing import Optional
from huggingface_hub import snapshot_download
from huggingface_hub.utils import validate_repo_id, HFValidationError
from ..llm import Config
from .llm import 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
class AutoModelForCausalLM:
@classmethod
def from_pretrained(
cls,
model_path_or_repo_id: str,
*,
local_files_only: bool = False,
revision: Optional[str] = None,
**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.
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.
Returns:
`LLM` object.
"""
config = Config()
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)
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.")
model_path = None
if path_type == "file":
model_path = Path(model_path_or_repo_id).parent
elif path_type == "dir":
model_path = Path(model_path_or_repo_id)
elif path_type == "repo":
model_path = snapshot_download(
repo_id=model_path_or_repo_id,
local_files_only=local_files_only,
revision=revision,
)
return LLM(model_path=model_path, config=config)