Make use of trust_remote_code consistent (#1806)

Co-authored-by: Quentin Gallouédec <quentin.gallouedec@huggingface.co>
huggingface · Jul 10, 2024 · 7ddef5c · 7ddef5c
1 parent a9cddf8
commit 7ddef5c
Show file tree

Hide file tree

Showing 13 changed files with 113 additions and 45 deletions.
diff --git a/examples/scripts/bco.py b/examples/scripts/bco.py
@@ -166,10 +166,16 @@ def mean_pooling(model_output, attention_mask):
     kto_args.gradient_checkpointing_kwargs = {"use_reentrant": True}
 
     # Load a pretrained model
-    model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path)
-    model_ref = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
+    )
+    model_ref = AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
+    )
 
-    tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
+    )
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
@@ -193,13 +199,15 @@ def format_dataset(example):
     accelerator = Accelerator()
     embedding_model = AutoModel.from_pretrained(
         "nomic-ai/nomic-embed-text-v1.5",
-        trust_remote_code=True,
+        trust_remote_code=model_args.trust_remote_code,
         safe_serialization=True,
         torch_dtype=torch.bfloat16,
         device_map="auto",
     )
     embedding_model = accelerator.prepare_model(embedding_model)
-    embedding_tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+    embedding_tokenizer = AutoTokenizer.from_pretrained(
+        "bert-base-uncased", trust_remote_code=model_args.trust_remote_code
+    )
     embedding_func = partial(
         embed_prompt,
         model=embedding_model,

diff --git a/examples/scripts/chat.py b/examples/scripts/chat.py
@@ -211,19 +211,24 @@ def parse_settings(user_input, current_args, interface):
 
 
 def load_model_and_tokenizer(args):
-    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, revision=args.model_revision)
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model_name_or_path,
+        revision=args.model_revision,
+        trust_remote_code=args.trust_remote_code,
+    )
 
     torch_dtype = args.torch_dtype if args.torch_dtype in ["auto", None] else getattr(torch, args.torch_dtype)
     quantization_config = get_quantization_config(args)
     model_kwargs = dict(
         revision=args.model_revision,
-        trust_remote_code=args.trust_remote_code,
         attn_implementation=args.attn_implementation,
         torch_dtype=torch_dtype,
         device_map="auto",
         quantization_config=quantization_config,
     )
-    model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path, **model_kwargs)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model_name_or_path, trust_remote_code=args.trust_remote_code, **model_kwargs
+    )
 
     if getattr(model, "hf_device_map", None) is None:
         model = model.to(args.device)

diff --git a/examples/scripts/cpo.py b/examples/scripts/cpo.py
@@ -76,8 +76,12 @@ class ScriptArguments:
     ################
     # Model & Tokenizer
     ################
-    model = AutoModelForCausalLM.from_pretrained(model_config.model_name_or_path)
-    tokenizer = AutoTokenizer.from_pretrained(model_config.model_name_or_path)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code
+    )
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 

diff --git a/examples/scripts/dpo.py b/examples/scripts/dpo.py
@@ -106,20 +106,25 @@
     quantization_config = get_quantization_config(model_config)
     model_kwargs = dict(
         revision=model_config.model_revision,
-        trust_remote_code=model_config.trust_remote_code,
         attn_implementation=model_config.attn_implementation,
         torch_dtype=torch_dtype,
         use_cache=False if training_args.gradient_checkpointing else True,
         device_map=get_kbit_device_map() if quantization_config is not None else None,
         quantization_config=quantization_config,
     )
-    model = AutoModelForCausalLM.from_pretrained(model_config.model_name_or_path, **model_kwargs)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs
+    )
     peft_config = get_peft_config(model_config)
     if peft_config is None:
-        model_ref = AutoModelForCausalLM.from_pretrained(model_config.model_name_or_path, **model_kwargs)
+        model_ref = AutoModelForCausalLM.from_pretrained(
+            model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs
+        )
     else:
         model_ref = None
-    tokenizer = AutoTokenizer.from_pretrained(model_config.model_name_or_path)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code
+    )
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     if tokenizer.chat_template is None:

diff --git a/examples/scripts/kto.py b/examples/scripts/kto.py
@@ -76,10 +76,16 @@ class ScriptArguments:
     script_args, kto_args, model_args = parser.parse_args_into_dataclasses()
 
     # Load a pretrained model
-    model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path)
-    model_ref = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
+    )
+    model_ref = AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
+    )
 
-    tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
+    )
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 

diff --git a/examples/scripts/orpo.py b/examples/scripts/orpo.py
@@ -76,8 +76,12 @@ class ScriptArguments:
     ################
     # Model & Tokenizer
     ################
-    model = AutoModelForCausalLM.from_pretrained(model_config.model_name_or_path)
-    tokenizer = AutoTokenizer.from_pretrained(model_config.model_name_or_path)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code
+    )
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 

diff --git a/examples/scripts/ppo/ppo.py b/examples/scripts/ppo/ppo.py
@@ -53,15 +53,23 @@
     tokenizer = AutoTokenizer.from_pretrained(
         model_config.model_name_or_path,
         padding_side="left",
-        trust_remote_code=True,
+        trust_remote_code=model_config.trust_remote_code,
     )
     tokenizer.add_special_tokens({"pad_token": "[PAD]"})
     if tokenizer.chat_template is None:
         tokenizer.chat_template = SIMPLE_QUERY_CHAT_TEMPLATE
-    value_model = AutoModelForSequenceClassification.from_pretrained(config.reward_model_path, num_labels=1)
-    reward_model = AutoModelForSequenceClassification.from_pretrained(config.reward_model_path, num_labels=1)
-    ref_policy = AutoModelForCausalLM.from_pretrained(config.sft_model_path)
-    policy = AutoModelForCausalLM.from_pretrained(config.sft_model_path)
+    value_model = AutoModelForSequenceClassification.from_pretrained(
+        config.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1
+    )
+    reward_model = AutoModelForSequenceClassification.from_pretrained(
+        config.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1
+    )
+    ref_policy = AutoModelForCausalLM.from_pretrained(
+        config.sft_model_path, trust_remote_code=model_config.trust_remote_code
+    )
+    policy = AutoModelForCausalLM.from_pretrained(
+        config.sft_model_path, trust_remote_code=model_config.trust_remote_code
+    )
     ################
     # Dataset
     ################

diff --git a/examples/scripts/ppo/ppo_tldr.py b/examples/scripts/ppo/ppo_tldr.py
@@ -57,15 +57,23 @@
     tokenizer = AutoTokenizer.from_pretrained(
         model_config.model_name_or_path,
         padding_side="left",
-        trust_remote_code=True,
+        trust_remote_code=model_config.trust_remote_code,
     )
     tokenizer.add_special_tokens({"pad_token": "[PAD]"})
     if tokenizer.chat_template is None:
         tokenizer.chat_template = SIMPLE_QUERY_CHAT_TEMPLATE
-    value_model = AutoModelForSequenceClassification.from_pretrained(config.reward_model_path, num_labels=1)
-    reward_model = AutoModelForSequenceClassification.from_pretrained(config.reward_model_path, num_labels=1)
-    ref_policy = AutoModelForCausalLM.from_pretrained(config.sft_model_path)
-    policy = AutoModelForCausalLM.from_pretrained(config.sft_model_path)
+    value_model = AutoModelForSequenceClassification.from_pretrained(
+        config.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1
+    )
+    reward_model = AutoModelForSequenceClassification.from_pretrained(
+        config.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1
+    )
+    ref_policy = AutoModelForCausalLM.from_pretrained(
+        config.sft_model_path, trust_remote_code=model_config.trust_remote_code
+    )
+    policy = AutoModelForCausalLM.from_pretrained(
+        config.sft_model_path, trust_remote_code=model_config.trust_remote_code
+    )
     ################
     # Dataset
     ################

diff --git a/examples/scripts/reward_modeling.py b/examples/scripts/reward_modeling.py
@@ -57,13 +57,14 @@
     quantization_config = get_quantization_config(model_config)
     model_kwargs = dict(
         revision=model_config.model_revision,
-        trust_remote_code=model_config.trust_remote_code,
         device_map=get_kbit_device_map() if quantization_config is not None else None,
         quantization_config=quantization_config,
     )
-    tokenizer = AutoTokenizer.from_pretrained(model_config.model_name_or_path, use_fast=True)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, use_fast=True
+    )
     model = AutoModelForSequenceClassification.from_pretrained(
-        model_config.model_name_or_path, num_labels=1, **model_kwargs
+        model_config.model_name_or_path, num_labels=1, trust_remote_code=model_config.trust_remote_code, **model_kwargs
     )
 
     if model_config.lora_task_type != "SEQ_CLS":

diff --git a/examples/scripts/rloo/rloo.py b/examples/scripts/rloo/rloo.py
@@ -55,14 +55,20 @@
     tokenizer = AutoTokenizer.from_pretrained(
         model_config.model_name_or_path,
         padding_side="left",
-        trust_remote_code=True,
+        trust_remote_code=model_config.trust_remote_code,
     )
     tokenizer.add_special_tokens({"pad_token": "[PAD]"})
     if tokenizer.chat_template is None:
         tokenizer.chat_template = SIMPLE_QUERY_CHAT_TEMPLATE
-    reward_model = AutoModelForSequenceClassification.from_pretrained(config.reward_model_path, num_labels=1)
-    ref_policy = AutoModelForCausalLM.from_pretrained(config.sft_model_path)
-    policy = AutoModelForCausalLM.from_pretrained(config.sft_model_path)
+    reward_model = AutoModelForSequenceClassification.from_pretrained(
+        config.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1
+    )
+    ref_policy = AutoModelForCausalLM.from_pretrained(
+        config.sft_model_path, trust_remote_code=model_config.trust_remote_code
+    )
+    policy = AutoModelForCausalLM.from_pretrained(
+        config.sft_model_path, trust_remote_code=model_config.trust_remote_code
+    )
     ################
     # Dataset
     ################

diff --git a/examples/scripts/rloo/rloo_tldr.py b/examples/scripts/rloo/rloo_tldr.py
@@ -59,14 +59,20 @@
     tokenizer = AutoTokenizer.from_pretrained(
         model_config.model_name_or_path,
         padding_side="left",
-        trust_remote_code=True,
+        trust_remote_code=model_config.trust_remote_code,
     )
     tokenizer.add_special_tokens({"pad_token": "[PAD]"})
     if tokenizer.chat_template is None:
         tokenizer.chat_template = SIMPLE_QUERY_CHAT_TEMPLATE
-    reward_model = AutoModelForSequenceClassification.from_pretrained(config.reward_model_path, num_labels=1)
-    ref_policy = AutoModelForCausalLM.from_pretrained(config.sft_model_path)
-    policy = AutoModelForCausalLM.from_pretrained(config.sft_model_path)
+    reward_model = AutoModelForSequenceClassification.from_pretrained(
+        config.reward_model_path, trust_remote_code=model_config.trust_remote_code, num_labels=1
+    )
+    ref_policy = AutoModelForCausalLM.from_pretrained(
+        config.sft_model_path, trust_remote_code=model_config.trust_remote_code
+    )
+    policy = AutoModelForCausalLM.from_pretrained(
+        config.sft_model_path, trust_remote_code=model_config.trust_remote_code
+    )
     ################
     # Dataset
     ################

diff --git a/examples/scripts/sft.py b/examples/scripts/sft.py
@@ -108,7 +108,9 @@
         quantization_config=quantization_config,
     )
     training_args.model_init_kwargs = model_kwargs
-    tokenizer = AutoTokenizer.from_pretrained(model_config.model_name_or_path, use_fast=True)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, use_fast=True
+    )
     tokenizer.pad_token = tokenizer.eos_token
 
     ################

diff --git a/examples/scripts/vsft_llava.py b/examples/scripts/vsft_llava.py
@@ -125,18 +125,23 @@
     quantization_config = get_quantization_config(model_config)
     model_kwargs = dict(
         revision=model_config.model_revision,
-        trust_remote_code=model_config.trust_remote_code,
         attn_implementation=model_config.attn_implementation,
         torch_dtype=torch_dtype,
         device_map=get_kbit_device_map() if quantization_config is not None else None,
         quantization_config=quantization_config,
     )
-    tokenizer = AutoTokenizer.from_pretrained(model_config.model_name_or_path, use_fast=True)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, use_fast=True
+    )
     tokenizer.chat_template = LLAVA_CHAT_TEMPLATE
-    processor = AutoProcessor.from_pretrained(model_config.model_name_or_path)
+    processor = AutoProcessor.from_pretrained(
+        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code
+    )
     processor.tokenizer = tokenizer
 
-    model = LlavaForConditionalGeneration.from_pretrained(model_config.model_name_or_path, **model_kwargs)
+    model = LlavaForConditionalGeneration.from_pretrained(
+        model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, **model_kwargs
+    )
 
     ################
     # Create a data collator to encode text and image pairs