Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix undefined variable in library name inference #1964

Merged
merged 1 commit into from
Jul 18, 2024
Merged

Fix undefined variable in library name inference #1964

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 10 additions & 7 deletions optimum/exporters/tasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,9 @@ def get_transformers_tasks_to_model_mapping(tasks_to_model_loader, framework="pt
for model_loader in model_loaders:
model_loader_class = getattr(auto_modeling_module, model_loader, None)
if model_loader_class is not None:
# we can just update the model_type to model_class mapping since we only need one either way
# we can just update the model_type to model_class mapping since
# we can only have one task->model_type->model_class either way
# e.g. we merge automatic-speech-recognition's SpeechSeq2Seq and CTC models
tasks_to_model_mapping[task_name].update(model_loader_class._model_mapping._model_mapping)

return tasks_to_model_mapping
Expand Down Expand Up @@ -1767,6 +1769,7 @@ def _infer_library_from_model_or_model_class(
model: Optional[Union["PreTrainedModel", "TFPreTrainedModel", "DiffusionPipeline"]] = None,
model_class: Optional[Type[Union["PreTrainedModel", "TFPreTrainedModel", "DiffusionPipeline"]]] = None,
):
inferred_library_name = None
if model is not None and model_class is not None:
raise ValueError("Either a model or a model class must be provided, but both were given here.")
if model is None and model_class is None:
Expand All @@ -1775,20 +1778,20 @@ def _infer_library_from_model_or_model_class(
target_class_module = model.__class__.__module__ if model is not None else model_class.__module__

if target_class_module.startswith("sentence_transformers"):
library_name = "sentence_transformers"
inferred_library_name = "sentence_transformers"
elif target_class_module.startswith("transformers"):
library_name = "transformers"
inferred_library_name = "transformers"
elif target_class_module.startswith("diffusers"):
library_name = "diffusers"
inferred_library_name = "diffusers"
elif target_class_module.startswith("timm"):
library_name = "timm"
inferred_library_name = "timm"

if library_name is None:
if inferred_library_name is None:
raise ValueError(
"The library name could not be automatically inferred. If using the command-line, please provide the argument --library {transformers,diffusers,timm,sentence_transformers}. Example: `--library diffusers`."
)

return library_name
return inferred_library_name

@classmethod
def _infer_library_from_model_name_or_path(
Expand Down
Loading