test

huggingface · Dec 15, 2024 · fbff767 · fbff767
1 parent 6fb84c6
commit fbff767
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/optimum_benchmark/backends/py_txi/backend.py b/optimum_benchmark/backends/py_txi/backend.py
@@ -93,10 +93,10 @@ def create_no_weights_model(self) -> None:
             self.generation_config.save_pretrained(save_directory=self.no_weights_model)
 
     def load_model_with_no_weights(self) -> None:
-        original_volumes, self.config.volumes = self.config.volumes, {self.tmpdir.name: {"bind": "/data", "mode": "rw"}}
-        original_model, self.config.model = self.config.model, "/data/no_weights_model/"
+        self.config.volumes[self.tmpdir.name] = {"bind": "/no_weights_data/", "mode": "rw"}
+        original_model, self.config.model = self.config.model, "/no_weights_data/no_weights_model/"
         self.load_model_from_pretrained()
-        self.config.model, self.config.volumes = original_model, original_volumes
+        self.config.model, self.config.volumes = original_model
 
     def load_model_from_pretrained(self) -> None:
         if self.config.task in TEXT_GENERATION_TASKS:

diff --git a/optimum_benchmark/backends/py_txi/config.py b/optimum_benchmark/backends/py_txi/config.py
@@ -37,7 +37,7 @@ class PyTXIConfig(BackendConfig):
         metadata={"help": "Dictionary of volumes to mount inside the container."},
     )
     environment: List[str] = field(
-        default_factory=lambda: ["HUGGING_FACE_HUB_TOKEN"],
+        default_factory=lambda: ["HF_TOKEN"],
         metadata={"help": "List of environment variables to forward to the container from the host."},
     )