Prevent crash if id2label is None

This may be the case on old models
huggingface · Dec 4, 2023 · 0f460a2 · 0f460a2
1 parent f4cfc3e
commit 0f460a2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/docker_images/setfit/app/pipelines/text_classification.py b/docker_images/setfit/app/pipelines/text_classification.py
@@ -23,7 +23,7 @@ def __call__(self, inputs: str) -> List[Dict[str, float]]:
         """
         probs = self.model.predict_proba([inputs], as_numpy=True)
         if probs.ndim == 2:
-            id2label = getattr(self.model, "id2label", {})
+            id2label = getattr(self.model, "id2label", {}) or {}
             return [
                 [
                     {"label": id2label.get(idx, idx), "score": prob}