support hf_access_token

README.md

@@ -78,11 +78,12 @@ setup: |
   echo "Pre-downloading base models so that you won't have to wait for long once the app is ready..."
   python llm_tuner/download_base_model.py --base_model_names='decapoda-research/llama-7b-hf,nomic-ai/gpt4all-j'
 
-# Start the app. `wandb_api_key` and `wandb_project_name` are optional.
+# Start the app. `hf_access_token`, `wandb_api_key` and `wandb_project_name` are optional.
 run: |
   conda activate llm-tuner
   python llm_tuner/app.py \
     --data_dir='/data' \
+    --hf_access_token="$([ -f /data/secrets/hf_access_token.txt ] && cat /data/secrets/hf_access_token.txt | tr -d '\n')" \
     --wandb_api_key="$([ -f /data/secrets/wandb_api_key.txt ] && cat /data/secrets/wandb_api_key.txt | tr -d '\n')" \
     --wandb_project_name='llm-tuner' \
     --timezone='Atlantic/Reykjavik' \

app.py

@@ -29,6 +29,7 @@ def main(
     ui_dev_mode: Union[bool, None] = None,
     wandb_api_key: Union[str, None] = None,
     wandb_project: Union[str, None] = None,
+    hf_access_token: Union[str, None] = None,
     timezone: Union[str, None] = None,
     config: Union[str, None] = None,
 ):
@@ -45,6 +46,8 @@ def main(
 
     :param wandb_api_key: The API key for Weights & Biases. Setting either this or `wandb_project` will enable Weights & Biases.
     :param wandb_project: The default project name for Weights & Biases. Setting either this or `wandb_api_key` will enable Weights & Biases.
+
+    :param hf_access_token: Provide an access token to load private models form Hugging Face Hub. An access token can be created at https://huggingface.co/settings/tokens.
     '''
 
     config_from_file = read_yaml_config(config_path=config)

llama_lora/config.py

@@ -8,19 +8,25 @@ class Config:
     Stores the application configuration. This is a singleton class.
     """
 
+    # Where data is stored
     data_dir: str = ""
-    load_8bit: bool = False
 
+    # Model Related
     default_base_model_name: str = ""
     base_model_choices: Union[List[str], str] = []
-
+    load_8bit: bool = False
     trust_remote_code: bool = False
 
+    # Application Settings
     timezone: Any = pytz.UTC
 
+    # Authentication
     auth_username: Union[str, None] = None
     auth_password: Union[str, None] = None
 
+    # Hugging Face
+    hf_access_token: Union[str, None] = None
+
     # WandB
     enable_wandb: Union[bool, None] = None
     wandb_api_key: Union[str, None] = None

llama_lora/lib/finetune.py

@@ -71,6 +71,7 @@ def train(
     wandb_watch: str = "false",  # options: false | gradients | all
     wandb_log_model: str = "true",  # options: false | true
     additional_wandb_config: Union[dict, None] = None,
+    hf_access_token: Union[str, None] = None,
     status_message_callback: Any = None,
     params_info_callback: Any = None,
 ):
@@ -88,9 +89,11 @@ def train(
         additional_training_arguments = None
     if isinstance(additional_training_arguments, str):
         try:
-            additional_training_arguments = json.loads(additional_training_arguments)
+            additional_training_arguments = json.loads(
+                additional_training_arguments)
         except Exception as e:
-            raise ValueError(f"Could not parse additional_training_arguments: {e}")
+            raise ValueError(
+                f"Could not parse additional_training_arguments: {e}")
 
     if isinstance(additional_lora_config, str):
         additional_lora_config = additional_lora_config.strip()
@@ -183,11 +186,13 @@ def train(
 
     if status_message_callback:
         if isinstance(base_model, str):
-            cb_result = status_message_callback(f"Preparing model '{base_model}' for training...")
+            cb_result = status_message_callback(
+                f"Preparing model '{base_model}' for training...")
             if cb_result:
                 return
         else:
-            cb_result = status_message_callback("Preparing model for training...")
+            cb_result = status_message_callback(
+                "Preparing model for training...")
             if cb_result:
                 return
 
@@ -201,6 +206,7 @@ def train(
             torch_dtype=torch.float16,
             llm_int8_skip_modules=lora_modules_to_save,
             device_map=device_map,
+            use_auth_token=hf_access_token
         )
         if re.match("[^/]+/llama", model_name):
             print(f"Setting special tokens for LLaMA model {model_name}...")
@@ -213,11 +219,14 @@ def train(
     if isinstance(tokenizer, str):
         tokenizer_name = tokenizer
         try:
-            tokenizer = AutoTokenizer.from_pretrained(tokenizer)
+            tokenizer = AutoTokenizer.from_pretrained(
+                tokenizer, use_auth_token=hf_access_token
+            )
         except Exception as e:
             if 'LLaMATokenizer' in str(e):
                 tokenizer = LlamaTokenizer.from_pretrained(
                     tokenizer_name,
+                    use_auth_token=hf_access_token
                 )
             else:
                 raise e
@@ -243,7 +252,8 @@ def train(
             f"Got error while running prepare_model_for_int8_training(model), maybe the model has already be prepared. Original error: {e}.")
 
     if status_message_callback:
-        cb_result = status_message_callback("Preparing PEFT model for training...")
+        cb_result = status_message_callback(
+            "Preparing PEFT model for training...")
         if cb_result:
             return
 
@@ -299,7 +309,8 @@ def train(
         wandb.config.update({"model": {"all_params": all_params, "trainable_params": trainable_params,
                             "trainable%": 100 * trainable_params / all_params}})
     if params_info_callback:
-        cb_result = params_info_callback(all_params=all_params, trainable_params=trainable_params)
+        cb_result = params_info_callback(
+            all_params=all_params, trainable_params=trainable_params)
         if cb_result:
             return
 

llama_lora/models.py

@@ -47,7 +47,11 @@ def get_new_base_model(base_model_name):
     while True:
         try:
             model = _get_model_from_pretrained(
-                model_class, base_model_name, from_tf=from_tf, force_download=force_download)
+                model_class,
+                base_model_name,
+                from_tf=from_tf,
+                force_download=force_download
+            )
             break
         except Exception as e:
             if 'from_tf' in str(e):
@@ -83,7 +87,9 @@ def get_new_base_model(base_model_name):
     return model
 
 
-def _get_model_from_pretrained(model_class, model_name, from_tf=False, force_download=False):
+def _get_model_from_pretrained(
+        model_class, model_name,
+        from_tf=False, force_download=False):
     torch = get_torch()
     device = get_device()
 
@@ -97,7 +103,8 @@ def _get_model_from_pretrained(model_class, model_name, from_tf=False, force_dow
             device_map={'': 0},
             from_tf=from_tf,
             force_download=force_download,
-            trust_remote_code=Config.trust_remote_code
+            trust_remote_code=Config.trust_remote_code,
+            use_auth_token=Config.hf_access_token
         )
     elif device == "mps":
         return model_class.from_pretrained(
@@ -106,7 +113,8 @@ def _get_model_from_pretrained(model_class, model_name, from_tf=False, force_dow
             torch_dtype=torch.float16,
             from_tf=from_tf,
             force_download=force_download,
-            trust_remote_code=Config.trust_remote_code
+            trust_remote_code=Config.trust_remote_code,
+            use_auth_token=Config.hf_access_token
         )
     else:
         return model_class.from_pretrained(
@@ -115,7 +123,8 @@ def _get_model_from_pretrained(model_class, model_name, from_tf=False, force_dow
             low_cpu_mem_usage=True,
             from_tf=from_tf,
             force_download=force_download,
-            trust_remote_code=Config.trust_remote_code
+            trust_remote_code=Config.trust_remote_code,
+            use_auth_token=Config.hf_access_token
         )
 
 
@@ -133,13 +142,15 @@ def get_tokenizer(base_model_name):
     try:
         tokenizer = AutoTokenizer.from_pretrained(
             base_model_name,
-            trust_remote_code=Config.trust_remote_code
+            trust_remote_code=Config.trust_remote_code,
+            use_auth_token=Config.hf_access_token
         )
     except Exception as e:
         if 'LLaMATokenizer' in str(e):
             tokenizer = LlamaTokenizer.from_pretrained(
                 base_model_name,
-                trust_remote_code=Config.trust_remote_code
+                trust_remote_code=Config.trust_remote_code,
+                use_auth_token=Config.hf_access_token
             )
         else:
             raise e
@@ -210,6 +221,7 @@ def get_model(
                 torch_dtype=torch.float16,
                 # ? https://github.com/tloen/alpaca-lora/issues/21
                 device_map={'': 0},
+                use_auth_token=Config.hf_access_token
             )
         elif device == "mps":
             model = PeftModel.from_pretrained(
@@ -217,12 +229,14 @@ def get_model(
                 peft_model_name_or_path,
                 device_map={"": device},
                 torch_dtype=torch.float16,
+                use_auth_token=Config.hf_access_token
             )
         else:
             model = PeftModel.from_pretrained(
                 model,
                 peft_model_name_or_path,
                 device_map={"": device},
+                use_auth_token=Config.hf_access_token
             )
 
     if re.match("[^/]+/llama", base_model_name):