fix: Back to InferenceClient with correct method signatures + diagnostics

src/hf_inference.py

@@ -1,8 +1,8 @@
 """
 HuggingFace Inference API client for Bayan models.
 
-Uses the `requests` library (via huggingface_hub's internal session)
-to call the HF Inference API from within HF Spaces.
+Uses huggingface_hub.InferenceClient which routes through HF's internal
+network when running inside HF Spaces (bypasses external DNS).
 
 Models:
   - bayan10/summarization-model  (MBart, summarization pipeline)
@@ -15,51 +15,23 @@ import os
 import json
 import logging
 import time
-import requests
 
 logger = logging.getLogger(__name__)
 
 HF_API_TOKEN = os.environ.get("HF_API_TOKEN", "")
-HF_API_BASE = "https://api-inference.huggingface.co/models/"
-HF_TIMEOUT = 120  # seconds — accounts for cold starts
 
+# Lazy-initialized client
+_client = None
 
-def _headers():
-    """Build request headers with auth token."""
-    h = {"Content-Type": "application/json"}
-    if HF_API_TOKEN:
-        h["Authorization"] = "Bearer " + HF_API_TOKEN
-    return h
 
-
-def _call_model(repo_id, payload):
-    """
-    Call any HF model via the Inference API.
-    Uses requests library + wait_for_model option.
-    Returns parsed JSON response.
-    """
-    url = HF_API_BASE + repo_id
-
-    # Tell HF to wait for the model to load instead of returning 503
-    if "options" not in payload:
-        payload["options"] = {"wait_for_model": True}
-
-    logger.info("HF API call: %s", repo_id)
-    resp = requests.post(url, headers=_headers(), json=payload, timeout=HF_TIMEOUT)
-
-    logger.info("HF API response for %s: HTTP %d, size=%d bytes",
-                repo_id, resp.status_code, len(resp.content))
-
-    if resp.status_code != 200:
-        logger.error("HF API error for %s: HTTP %d — %s",
-                     repo_id, resp.status_code, resp.text[:500])
-        raise RuntimeError("HF API error for {} (HTTP {}): {}".format(
-            repo_id, resp.status_code, resp.text[:300]))
-
-    result = resp.json()
-    logger.info("HF API result for %s: type=%s, preview=%s",
-                repo_id, type(result).__name__, str(result)[:150])
-    return result
+def _get_client():
+    """Get or create the InferenceClient singleton."""
+    global _client
+    if _client is None:
+        from huggingface_hub import InferenceClient
+        _client = InferenceClient(token=HF_API_TOKEN if HF_API_TOKEN else None)
+        logger.info("InferenceClient initialized (token=%s)", "set" if HF_API_TOKEN else "not set")
+    return _client
 
 
 # ============================================================
@@ -73,83 +45,104 @@ AUTOCOMPLETE_REPO = os.environ.get("AUTOCOMPLETE_REPO_ID", "bayan10/AutoComplete
 
 
 # ============================================================
-# Model-specific wrappers
+# Model-specific wrappers using InferenceClient typed methods
 # ============================================================
 
-def _extract_text(result, fallback=""):
-    """Extract generated text from various HF response formats."""
-    if isinstance(result, list) and len(result) > 0:
-        item = result[0]
-        if isinstance(item, dict):
-            return (item.get("summary_text")
-                    or item.get("generated_text")
-                    or item.get("translation_text")
-                    or fallback)
-        return str(item) if str(item).strip() else fallback
-
-    if isinstance(result, dict):
-        return (result.get("summary_text")
-                or result.get("generated_text")
-                or result.get("translation_text")
-                or fallback)
+def hf_summarize(text, max_length=128, min_length=30):
+    """Summarize Arabic text via HF Inference API."""
+    client = _get_client()
+    logger.info("Calling summarization: %s", SUMMARIZATION_REPO)
 
-    return str(result) if result else fallback
+    result = client.summarization(text, model=SUMMARIZATION_REPO)
 
+    logger.info("Summarization result: %s — %s", type(result).__name__, str(result)[:150])
 
-def hf_summarize(text, max_length=128, min_length=30):
-    """Summarize Arabic text via HF Inference API."""
-    result = _call_model(SUMMARIZATION_REPO, {
-        "inputs": text,
-        "parameters": {
-            "max_length": max_length,
-            "min_length": min_length,
-        }
-    })
-    return _extract_text(result, text[:100])
+    # SummarizationOutput has .summary_text
+    if hasattr(result, "summary_text"):
+        return result.summary_text
+    if isinstance(result, dict):
+        return result.get("summary_text", result.get("generated_text", str(result)))
+    return str(result)
 
 
 def hf_correct_spelling(text):
     """Correct spelling in Arabic text via HF Inference API."""
-    result = _call_model(SPELLING_REPO, {"inputs": text})
-    return _extract_text(result, text)
+    client = _get_client()
+    logger.info("Calling spelling: %s", SPELLING_REPO)
+
+    # Try text2text_generation first (for seq2seq models), fall back to text_generation
+    try:
+        result = client.text2text_generation(text, model=SPELLING_REPO)
+        logger.info("Spelling result (t2t): %s — %s", type(result).__name__, str(result)[:150])
+        if hasattr(result, "generated_text"):
+            return result.generated_text
+        if isinstance(result, str):
+            return result if result.strip() else text
+        if isinstance(result, dict):
+            return result.get("generated_text", text)
+        return text
+    except Exception as e1:
+        logger.warning("text2text_generation failed for spelling: %s", repr(e1)[:200])
+        try:
+            result = client.text_generation(text, model=SPELLING_REPO, max_new_tokens=len(text) + 50)
+            logger.info("Spelling result (tg): %s — %s", type(result).__name__, str(result)[:150])
+            if isinstance(result, str):
+                return result if result.strip() else text
+            return text
+        except Exception as e2:
+            logger.error("text_generation also failed for spelling: %s", repr(e2)[:200])
+            raise
 
 
 def hf_add_punctuation(text):
     """Add punctuation to Arabic text via HF Inference API."""
-    result = _call_model(PUNCTUATION_REPO, {"inputs": text})
-    return _extract_text(result, text)
+    client = _get_client()
+    logger.info("Calling punctuation: %s", PUNCTUATION_REPO)
+
+    try:
+        result = client.text2text_generation(text, model=PUNCTUATION_REPO)
+        logger.info("Punctuation result (t2t): %s — %s", type(result).__name__, str(result)[:150])
+        if hasattr(result, "generated_text"):
+            return result.generated_text
+        if isinstance(result, str):
+            return result if result.strip() else text
+        if isinstance(result, dict):
+            return result.get("generated_text", text)
+        return text
+    except Exception as e1:
+        logger.warning("text2text_generation failed for punctuation: %s", repr(e1)[:200])
+        try:
+            result = client.text_generation(text, model=PUNCTUATION_REPO, max_new_tokens=len(text) + 50)
+            logger.info("Punctuation result (tg): %s — %s", type(result).__name__, str(result)[:150])
+            if isinstance(result, str):
+                return result if result.strip() else text
+            return text
+        except Exception as e2:
+            logger.error("text_generation also failed for punctuation: %s", repr(e2)[:200])
+            raise
 
 
 def hf_autocomplete(text, n=5):
     """Get autocomplete suggestions for Arabic text via HF Inference API."""
-    result = _call_model(AUTOCOMPLETE_REPO, {
-        "inputs": text,
-        "parameters": {"max_new_tokens": 20}
-    })
+    client = _get_client()
+    logger.info("Calling autocomplete: %s", AUTOCOMPLETE_REPO)
+
+    result = client.text_generation(text, model=AUTOCOMPLETE_REPO, max_new_tokens=20)
+
+    logger.info("Autocomplete result: %s — %s", type(result).__name__, str(result)[:150])
 
     if isinstance(result, str):
         completion = result[len(text):].strip() if result.startswith(text) else result
         return [completion] if completion else [text]
 
-    if isinstance(result, list):
-        suggestions = []
-        for item in result:
-            gen = item.get("generated_text", "") if isinstance(item, dict) else str(item)
-            if gen.startswith(text):
-                gen = gen[len(text):].strip()
-            if gen:
-                suggestions.append(gen)
-        return suggestions if suggestions else [text]
-
     return [text]
 
 
 def check_hf_api_available():
     """Quick check if HF Inference API is reachable."""
     try:
-        resp = requests.get(HF_API_BASE + SUMMARIZATION_REPO,
-                           headers=_headers(), timeout=10)
-        return resp.status_code == 200
+        client = _get_client()
+        return client is not None
     except Exception:
         return False
 
@@ -157,12 +150,38 @@ def check_hf_api_available():
 def debug_test_all_models():
     """
     Test all HF models and return results dict.
-    Used by /api/debug-models endpoint for troubleshooting.
+    Also includes diagnostic info about InferenceClient internals.
     """
     results = {}
     test_text = "هذا نص تجريبي للاختبار"
     long_text = (test_text + " ") * 5
 
+    # Diagnostic info
+    try:
+        client = _get_client()
+        diag = {
+            "client_type": type(client).__name__,
+            "api_url": getattr(client, "api_url", "N/A"),
+            "base_url": getattr(client, "base_url", "N/A"),
+            "model": getattr(client, "model", "N/A"),
+        }
+        # Check available methods
+        diag["has_post"] = hasattr(client, "post")
+        diag["has_text2text"] = hasattr(client, "text2text_generation")
+        diag["has_summarization"] = hasattr(client, "summarization")
+        diag["has_text_generation"] = hasattr(client, "text_generation")
+    except Exception as e:
+        diag = {"error": repr(e)[:200]}
+
+    results["_diagnostics"] = diag
+
+    # Test env vars
+    results["_env"] = {
+        "HF_INFERENCE_ENDPOINT": os.environ.get("HF_INFERENCE_ENDPOINT", "NOT SET"),
+        "HF_API_URL": os.environ.get("HF_API_URL", "NOT SET"),
+        "SPACE_ID": os.environ.get("SPACE_ID", "NOT SET"),
+    }
+
     for name, fn, args in [
         ("summarization", hf_summarize, (long_text, 30, 10)),
         ("spelling", hf_correct_spelling, (test_text,)),