From 0f460a2c8f5b0a3559e70fc061ef0413fb9fbdc3 Mon Sep 17 00:00:00 2001
From: Tom Aarsen <Cubiegamedev@gmail.com>
Date: Mon, 4 Dec 2023 09:25:17 +0100
Subject: [PATCH] Prevent crash if id2label is None

This may be the case on old models
---
 docker_images/setfit/app/pipelines/text_classification.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker_images/setfit/app/pipelines/text_classification.py b/docker_images/setfit/app/pipelines/text_classification.py
index 11d365bf..681c740a 100644
--- 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}