Update services/huggingface_api.py

services/huggingface_api.py

@@ -1,96 +1,55 @@
 import os
-import time
 import logging
-import requests
-from openai import OpenAI
+from huggingface_hub import InferenceClient
 import numpy as np
 
 logger = logging.getLogger(__name__)
 
-HF_API_KEY = os.getenv('HF_API_KEY')
+HF_API_KEY = os.getenv("HF_API_KEY")
 if not HF_API_KEY:
     raise ValueError("HF_API_KEY not set")
 
-HF_HEADERS = {"Authorization": f"Bearer {HF_API_KEY}"}
-HF_BASE = "https://api-inference.huggingface.co"
+client = InferenceClient(
+    provider="hf-inference",
+    api_key=HF_API_KEY,
+)
 
-def hf_post(url, payload, retry=3):
-    for attempt in range(retry + 1):
-        try:
-            resp = requests.post(url, headers=HF_HEADERS, json=payload, timeout=60)
-            if resp.status_code == 200:
-                return resp.json()
-            else:
-                logger.warning(f"HF call to {url} status {resp.status_code} attempt {attempt}: {resp.text}")
-                time.sleep(2 ** attempt)
-        except requests.exceptions.RequestException as e:
-            logger.warning(f"HF request to {url} failed attempt {attempt}: {e}")
-            time.sleep(2 ** attempt)
-    raise Exception(f"Hugging Face API failed for {url}: {resp.status_code} {resp.text}")
-
-# -------------------------
-# Embeddings
-# -------------------------
 def get_embeddings(model, text):
-    """
-    Calls HF feature-extraction to get embeddings for a single text.
-    Returns 1D list of floats.
-    """
-    url = f"{HF_BASE}/pipeline/feature-extraction/{model}"
-    payload = {"inputs": text, "options": {"wait_for_model": True}}
+    if not text or not isinstance(text, str):
+        logger.error(f"Invalid input for embeddings: {text}")
+        raise ValueError("Text must be a non-empty string")
     try:
-        out = hf_post(url, payload)
-        # Handle response: list of token embeddings for the input text
+        # Feature extraction returns a list of token embeddings
+        out = client.feature_extraction(text, model=model)
+        logger.debug(f"HF feature extraction response: {out}")
         if isinstance(out, list) and len(out) > 0:
-            candidate = out[0]
-            if isinstance(candidate, list) and len(candidate) > 0:
-                # Average token embeddings to get sentence embedding
-                arr = np.array(candidate)
-                vec = arr.mean(axis=0).tolist()
-                return vec
-            else:
-                raise Exception(f"Unexpected embeddings format: {candidate}")
+            arr = np.array(out[0])  # Take the first sequence's embeddings
+            vec = arr.mean(axis=0).tolist()  # Average token embeddings
+            return vec
         else:
-            raise Exception(f"Unexpected HF response: {out}")
+            raise Exception(f"Unexpected feature extraction response: {out}")
     except Exception as e:
         logger.exception(f"Failed to get embeddings for {model}: {e}")
         raise
 
-# -------------------------
-# Cross Encoder
-# -------------------------
 def get_cross_encoder_score(model, student, teacher):
-    """
-    Obtain a similarity/score between student and teacher.
-    Returns float between 0 and 1.
-    """
-    url = f"{HF_BASE}/pipeline/sentence-similarity/{model}"
-    payload = {"inputs": {"source_sentence": student, "sentences": [teacher]}}
     try:
-        out = hf_post(url, payload)
+        # Sentence similarity expects a dict with source_sentence and sentences
+        out = client.sentence_similarity(
+            {"source_sentence": student, "sentences": [teacher]},
+            model=model
+        )
+        logger.debug(f"HF sentence similarity response: {out}")
         if isinstance(out, list) and len(out) > 0 and isinstance(out[0], (float, int)):
             return float(out[0])
-        raise Exception(f"Unexpected cross-encoder response: {out}")
+        raise Exception(f"Unexpected sentence similarity response: {out}")
     except Exception as e:
         logger.exception(f"Failed cross-encoder for {model}: {e}")
-        # Fallback to default model endpoint
-        url = f"{HF_BASE}/{model}"
-        payload = {"inputs": [student, teacher]}
-        out = hf_post(url, payload)
-        if isinstance(out, list) and len(out) > 0:
-            return float(out[0]) if isinstance(out[0], (float, int)) else float(out[0].get("score", 0.0))
-        raise Exception(f"Unable to parse cross-encoder output: {out}")
-
-# -------------------------
-# Feedback Generation
-# -------------------------
-# ... (unchanged)
+        raise
 
-# ------------------------- 
-# Feedback Generation 
-# ------------------------- 
 def generate_feedback(question, answer, score, model="meta-llama/Llama-3.2-1B-Instruct:novita"):
     try:
+        from openai import OpenAI
         client = OpenAI(
             base_url="https://router.huggingface.co/v1",
             api_key=os.getenv("HF_API_KEY"),