Use InferenceClient for HF chat

app.py

@@ -14,6 +14,8 @@ import streamlit as st
 import numpy as np
 import faiss
 import requests
+from huggingface_hub import InferenceClient
+from huggingface_hub import InferenceClient
 from sentence_transformers import SentenceTransformer
 
 load_dotenv(Path(__file__).resolve().parent / ".env", override=True)
@@ -37,9 +39,6 @@ HF_TOKEN = (
     or ""
 ).strip()
 HF_MODEL = os.environ.get("RAG_HF_MODEL", "meta-llama/Llama-3.2-1B-Instruct")
-HF_API_URL = os.environ.get("RAG_HF_API_URL", "").strip()
-if not HF_API_URL:
-    HF_API_URL = f"https://api-inference.huggingface.co/models/{HF_MODEL}"
 
 OLLAMA_BASE_URL = os.environ.get("RAG_OLLAMA_URL", "http://localhost:11434").rstrip("/")
 OLLAMA_MODEL = os.environ.get("RAG_OLLAMA_MODEL", "llama3.2:1b")
@@ -49,14 +48,14 @@ REPO_NAME = "RAG_project"
 
 GITHUB_TOKEN = os.environ.get("GITHUB_TOKEN", "").strip()
 # Retrieval mix: book-first + article nuance.
-BOOK_K = 12
+BOOK_K = 8
 ARTICLE_K = 4
 PER_DOC_CAP = 3
 OVERLAP_FILTER = True
 
 # Enhanced answer mix: heavier retrieval for deeper answers.
-ENHANCED_BOOK_K = 12
-ENHANCED_ARTICLE_K = 5
+ENHANCED_BOOK_K = 14
+ENHANCED_ARTICLE_K = 7
 
 AVOID_PHRASES = [
     "The article discusses",
@@ -523,31 +522,24 @@ def build_hf_prompt(user_prompt: str, model_id: str) -> str:
         )
     return f"System: {system_msg}\nUser: {user_prompt}\nAssistant:"
 
+@st.cache_resource(show_spinner=False)
+def get_hf_client() -> InferenceClient:
+    return InferenceClient(model=HF_MODEL, token=HF_TOKEN)
+
 def hf_chat(prompt: str, timeout: Tuple[int, int] = (10, 600)) -> Tuple[str, Optional[str]]:
     if not HF_TOKEN:
         return "", "Missing HF_TOKEN (or HUGGINGFACEHUB_API_TOKEN)"
-    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-    payload = {
-        "inputs": build_hf_prompt(prompt, HF_MODEL),
-        "parameters": {
-            "temperature": 0.2,
-            "max_new_tokens": 512,
-            "return_full_text": False,
-        },
-        "options": {"wait_for_model": True},
-    }
     try:
-        r = requests.post(HF_API_URL, headers=headers, json=payload, timeout=timeout)
-        r.raise_for_status()
-        data = r.json()
-        msg = ""
-        if isinstance(data, list) and data:
-            msg = data[0].get("generated_text") or ""
-        elif isinstance(data, dict):
-            if "error" in data:
-                return "", str(data.get("error") or "Hugging Face error")
-            msg = data.get("generated_text") or ""
-        return (msg or "").strip(), None
+        client = get_hf_client()
+        inp = build_hf_prompt(prompt, HF_MODEL)
+        out = client.text_generation(
+            inp,
+            max_new_tokens=512,
+            temperature=0.2,
+            do_sample=True,
+            return_full_text=False,
+        )
+        return (out or "").strip(), None
     except Exception as e:
         return "", str(e)
 

requirements.txt

@@ -5,6 +5,7 @@ requests>=2.31.0
 python-dateutil>=2.9.0
 tqdm>=4.66.0
 python-dotenv>=1.0.0
+huggingface-hub>=0.30.0
 
 # -------------------------
 # PDF processing (required by normalize_all.py)