support --trust_remote_code, resolves #6

app.py

@@ -15,6 +15,7 @@ def main(
     base_model: str = "",
     data_dir: str = "",
     base_model_choices: str = "",
+    trust_remote_code: bool = False,
     # Allows to listen on all interfaces by providing '0.0.0.0'.
     server_name: str = "127.0.0.1",
     share: bool = False,
@@ -60,6 +61,8 @@ def main(
     if base_model not in Global.base_model_choices:
         Global.base_model_choices = [base_model] + Global.base_model_choices
 
+    Global.trust_remote_code = trust_remote_code
+
     Global.data_dir = os.path.abspath(data_dir)
     Global.load_8bit = load_8bit
 

llama_lora/globals.py

@@ -20,6 +20,8 @@ class Global:
     base_model_name: str = ""
     base_model_choices: List[str] = []
 
+    trust_remote_code = False
+
     # Functions
     train_fn: Any = train
 

llama_lora/models.py

@@ -37,16 +37,21 @@ def get_new_base_model(base_model_name):
             # device_map="auto",
             # ? https://github.com/tloen/alpaca-lora/issues/21
             device_map={'': 0},
+            trust_remote_code=Global.trust_remote_code
         )
     elif device == "mps":
         model = AutoModelForCausalLM.from_pretrained(
             base_model_name,
             device_map={"": device},
             torch_dtype=torch.float16,
+            trust_remote_code=Global.trust_remote_code
         )
     else:
         model = AutoModelForCausalLM.from_pretrained(
-            base_model_name, device_map={"": device}, low_cpu_mem_usage=True
+            base_model_name,
+            device_map={"": device},
+            low_cpu_mem_usage=True,
+            trust_remote_code=Global.trust_remote_code
         )
 
     tokenizer = get_tokenizer(base_model_name)
@@ -68,10 +73,16 @@ def get_tokenizer(base_model_name):
         return loaded_tokenizer
 
     try:
-        tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+        tokenizer = AutoTokenizer.from_pretrained(
+            base_model_name,
+            trust_remote_code=Global.trust_remote_code
+        )
     except Exception as e:
         if 'LLaMATokenizer' in str(e):
-            tokenizer = LlamaTokenizer.from_pretrained(base_model_name)
+            tokenizer = LlamaTokenizer.from_pretrained(
+                base_model_name,
+                trust_remote_code=Global.trust_remote_code
+            )
         else:
             raise e
 
@@ -100,13 +111,15 @@ def get_model(
     peft_model_name_or_path = peft_model_name
 
     if peft_model_name:
-        lora_models_directory_path = os.path.join(Global.data_dir, "lora_models")
+        lora_models_directory_path = os.path.join(
+            Global.data_dir, "lora_models")
         possible_lora_model_path = os.path.join(
             lora_models_directory_path, peft_model_name)
         if os.path.isdir(possible_lora_model_path):
             peft_model_name_or_path = possible_lora_model_path
 
-            possible_model_info_json_path = os.path.join(possible_lora_model_path, "info.json")
+            possible_model_info_json_path = os.path.join(
+                possible_lora_model_path, "info.json")
             if os.path.isfile(possible_model_info_json_path):
                 try:
                     with open(possible_model_info_json_path, "r") as file:
@@ -115,7 +128,8 @@ def get_model(
                         if possible_hf_model_name and json_data.get("load_from_hf"):
                             peft_model_name_or_path = possible_hf_model_name
                 except Exception as e:
-                    raise ValueError("Error reading model info from {possible_model_info_json_path}: {e}")
+                    raise ValueError(
+                        "Error reading model info from {possible_model_info_json_path}: {e}")
 
     Global.loaded_models.prepare_to_set()
     clear_cache()
@@ -148,7 +162,8 @@ def get_model(
             )
 
     if re.match("[^/]+/llama", base_model_name):
-        model.config.pad_token_id = get_tokenizer(base_model_name).pad_token_id = 0
+        model.config.pad_token_id = get_tokenizer(
+            base_model_name).pad_token_id = 0
         model.config.bos_token_id = 1
         model.config.eos_token_id = 2
 
@@ -166,7 +181,8 @@ def get_model(
 
 
 def prepare_base_model(base_model_name=Global.default_base_model_name):
-    Global.new_base_model_that_is_ready_to_be_used = get_new_base_model(base_model_name)
+    Global.new_base_model_that_is_ready_to_be_used = get_new_base_model(
+        base_model_name)
     Global.name_of_new_base_model_that_is_ready_to_be_used = base_model_name