From eeba6b0a932772b5ebdba58b43f29a6cd90eb93d Mon Sep 17 00:00:00 2001 From: Yu Li Date: Wed, 20 Dec 2023 11:50:34 -0600 Subject: [PATCH] support hf_token in auto model --- air_llm/airllm/auto_model.py | 11 ++++++++--- air_llm/setup.py | 2 +- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/air_llm/airllm/auto_model.py b/air_llm/airllm/auto_model.py index a503a48..f013b94 100644 --- a/air_llm/airllm/auto_model.py +++ b/air_llm/airllm/auto_model.py @@ -16,8 +16,12 @@ class AutoModel: "using the `AutoModel.from_pretrained(pretrained_model_name_or_path)` method." ) @classmethod - def get_module_class(cls, pretrained_model_name_or_path): - config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True) + def get_module_class(cls, pretrained_model_name_or_path, *inputs, **kwargs): + if 'hf_token' in kwargs: + print(f"using hf_token") + config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True, token=kwargs['hf_token']) + else: + config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True) if "QWen" in config.architectures[0]: return "airllm", "AirLLMQWen" @@ -37,7 +41,8 @@ class AutoModel: @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs): - module, cls = AutoModel.get_module_class(pretrained_model_name_or_path) + + module, cls = AutoModel.get_module_class(pretrained_model_name_or_path, *inputs, **kwargs) module = importlib.import_module(module) class_ = getattr(module, cls) diff --git a/air_llm/setup.py b/air_llm/setup.py index 60e6387..4c61c92 100644 --- a/air_llm/setup.py +++ b/air_llm/setup.py @@ -5,7 +5,7 @@ with open("README.md", "r") as fh: setuptools.setup( name="airllm", - version="2.6.1", + version="2.6.2", author="Gavin Li", author_email="gavinli@animaai.cloud", description="AirLLM allows single 4GB GPU card to run 70B large language models without quantization, distillation or pruning.",