Fix tests.

ludwig-ai · Oct 27, 2023 · b7ac42c · b7ac42c
1 parent e42d3e3
commit b7ac42c
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/ludwig/models/llm.py b/ludwig/models/llm.py
@@ -102,7 +102,7 @@ def load_pretrained_from_config(
 
     logger.info("Loading large language model...")
     pretrained_model_name_or_path = weights_save_path or config_obj.base_model
-    model: PreTrainedModel = load_pretrained_hf_class_with_hub_fallback(
+    model, _ = load_pretrained_hf_class_with_hub_fallback(
         AutoModelForCausalLM, pretrained_model_name_or_path, **load_kwargs
     )
     return model
@@ -126,7 +126,7 @@ def __init__(
         self._random_seed = random_seed
 
         self.model_name = self.config_obj.base_model
-        self.model_config = load_pretrained_hf_class_with_hub_fallback(AutoConfig, self.config_obj.base_model)
+        self.model_config, _ = load_pretrained_hf_class_with_hub_fallback(AutoConfig, self.config_obj.base_model)
 
         self.model = load_pretrained_from_config(self.config_obj, model_config=self.model_config)
         self.curr_device = next(self.model.parameters()).device
@@ -303,7 +303,7 @@ def to_device(self, device):
                 if self.config_obj.adapter:
                     from peft import PeftModel
 
-                    self.model = load_pretrained_hf_class_with_hub_fallback(
+                    self.model, _ = load_pretrained_hf_class_with_hub_fallback(
                         AutoModelForCausalLM,
                         self.model_name,
                         **model_kwargs,
@@ -314,7 +314,7 @@ def to_device(self, device):
                         torch_dtype=torch.float16,
                     )
                 else:
-                    self.model = load_pretrained_hf_class_with_hub_fallback(
+                    self.model, _ = load_pretrained_hf_class_with_hub_fallback(
                         AutoModelForCausalLM,
                         tmpdir,
                         **model_kwargs,
@@ -698,7 +698,7 @@ def load(self, save_path):
                 # Unwrap and reload PeftModel
                 self.model = self.model.base_model
 
-            self.model = load_pretrained_hf_class_with_hub_fallback(PeftModel, self.model, weights_save_path)
+            self.model = PeftModel.from_pretrained(self.model, weights_save_path)
         elif self.config_obj.trainer.type != "none":
             self.model = load_pretrained_from_config(
                 self.config_obj, model_config=self.model_config, weights_save_path=weights_save_path