Update app.py

app.py

@@ -19,8 +19,9 @@ LABELS = [
 ALLOWED = set(LABELS)
 
 def log(msg, **kv):
-    print(" | ".join([msg] + [f"{k}={v}"]) for k, v in kv.items())
-    print(" | ".join([msg] + [f"{k}={v}" for k, v in kv.items()]), file=sys.stderr, flush=True)
+    line = " | ".join([msg] + [f"{k}={v}" for k, v in kv.items()])
+    print(line, file=sys.stderr, flush=True)
+
 
 # Init HF client once
 client = InferenceClient(token=API_KEY) if API_KEY else None
@@ -52,29 +53,42 @@ def generate_text():
         return jsonify({"response": "generate_image"}), 200
 
     try:
-        # Zero-shot classification
         log("DECISION_CALL_ZSL", model=ZSL_MODEL_ID, prompt_len=len(prompt))
         zs = client.zero_shot_classification(
             prompt,
             LABELS,
             model=ZSL_MODEL_ID,
             hypothesis_template="This text is about {}.",
-            multi_label=False,          # pick the best single label
+            multi_label=False,  # single best label
         )
-        # Response shape: {'labels': [...], 'scores': [...], 'sequence': '...'}
-        labels = zs.get("labels") or []
+    
+        # Normalize shapes:
+        # - Newer hub often returns a dict
+        # - Some providers return a list[dict] (one per input)
+        if isinstance(zs, list):
+            zs = zs[0] if zs else {}
+    
+        if not isinstance(zs, dict):
+            raise ValueError(f"Unexpected ZSL response type: {type(zs)}")
+    
+        labels = zs.get("labels") or zs.get("candidate_labels") or []
         scores = zs.get("scores") or []
+        if not labels and "label" in zs:
+            labels = [zs["label"]]
+            scores = [zs.get("score", 0.0)]
+    
         best = labels[0] if labels else "other_query"
         score = float(scores[0]) if scores else 0.0
-
         token = best if best in ALLOWED else "other_query"
+    
         log("DECISION_OK", token=token, top_label=best, score=round(score, 4))
         return jsonify({"response": token}), 200
-
+    
     except Exception as e:
         log("DECISION_FAIL", error=str(e))
         return jsonify({"response": "other_query", "error": str(e)}), 200
 
+
 if __name__ == "__main__":
     port = int(os.getenv("PORT", 7860))
     log("BOOT", port=port, zsl_model=ZSL_MODEL_ID, api_key_set=bool(API_KEY))