fix inference client bug

src/entailment.py

@@ -1,7 +1,9 @@
+import os
 import json
 import time
 from concurrent.futures import ThreadPoolExecutor
 
+import requests
 import numpy as np
 import torch
 
@@ -67,37 +69,52 @@ def _parse_classification(data):
     return _label_to_id(best.get("label", ""))
 
 
-def _api_classify_pair(client, model, premise, hypothesis):
-    """One serverless text-classification call for a (premise, hypothesis) pair."""
-    payload = {"inputs": {"text": premise, "text_pair": hypothesis}}
-    for attempt in range(4):
+_API_URL = "https://api-inference.huggingface.co/models/" + ENTAILMENT_API_MODEL
+_warned_once = {"done": False}
+
+
+def _api_classify_pair(token, premise, hypothesis):
+    """One serverless text-classification call for a (premise, hypothesis) pair.
+    Uses the raw Inference API with the sentence-pair payload."""
+    headers = {"Authorization": f"Bearer {token}"} if token else {}
+    payload = {
+        "inputs": {"text": premise, "text_pair": hypothesis},
+        "options": {"wait_for_model": True},
+    }
+    for attempt in range(5):
         try:
-            raw = client.post(json=payload, model=model, task="text-classification")
-            return _parse_classification(raw)
-        except Exception as e:
-            msg = str(e).lower()
-            # 503 = model loading (cold start); 429 = rate limit → wait & retry
-            if any(x in msg for x in ("503", "loading", "429", "rate limit", "too many requests")):
+            r = requests.post(_API_URL, headers=headers, json=payload, timeout=60)
+            if r.status_code in (429, 503):
                 time.sleep(min(2 ** attempt * 2, 20))
                 continue
-            print(f"[entailment-api] classify error: {e}", flush=True)
-            return 1
+            if r.status_code >= 400:
+                if not _warned_once["done"]:
+                    print(f"[entailment-api] HTTP {r.status_code}: {r.text[:300]}", flush=True)
+                    _warned_once["done"] = True
+                return 1
+            return _parse_classification(r.json())
+        except Exception as e:
+            if not _warned_once["done"]:
+                print(f"[entailment-api] request error: {e}", flush=True)
+                _warned_once["done"] = True
+            time.sleep(min(2 ** attempt, 8))
     return 1
 
 
-def api_check_implication(pairs, client):
+def api_check_implication(pairs, client=None):
     """
     Entailment via HF serverless Inference API (parallel requests).
     Returns list[int] in {0,1,2}, aligned with `pairs`.
     """
     if not pairs:
         return []
+    token = getattr(client, "token", None) or os.environ.get("HF_TOKEN")
     total = len(pairs)
     print(f"[entailment-api] Classifying {total} pairs via Inference API "
           f"(model={ENTAILMENT_API_MODEL}, workers={ENTAILMENT_API_WORKERS})...", flush=True)
 
     def work(pair):
-        return _api_classify_pair(client, ENTAILMENT_API_MODEL, pair[0], pair[1])
+        return _api_classify_pair(token, pair[0], pair[1])
 
     with ThreadPoolExecutor(max_workers=ENTAILMENT_API_WORKERS) as ex:
         results = list(ex.map(work, pairs))