Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import requests
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
@@ -6,7 +7,6 @@ from langchain_community.document_loaders import PyPDFLoader
 from langchain_core.prompts import PromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough
-from huggingface_hub import InferenceClient
 
 import gradio as gr
 import warnings
@@ -53,11 +53,14 @@ def get_huggingface_token():
 # LLM
 # ---------------------------
 def get_llm(model_id: str = MODEL_OPTIONS[0], max_tokens: int = 256, temperature: float = 0.8):
+    """
+    Returns API URL, headers, and parameters for HuggingFace Inference API.
+    """
     token = get_huggingface_token()
+    api_url = f"https://api-inference.huggingface.co/models/{model_id}"
+    headers = {"Authorization": f"Bearer {token}"}
     
-    # Use InferenceClient directly for better reliability
-    client = InferenceClient(model=model_id, token=token)
-    return client, max_tokens, temperature
+    return api_url, headers, max_tokens, temperature
 
 
 # ---------------------------
@@ -140,7 +143,7 @@ def retriever_qa(file, query, model_choice, max_tokens, temperature, embedding_m
     
     try:
         selected_model = model_choice or MODEL_OPTIONS[0]
-        client, max_tok, temp = get_llm(selected_model, int(max_tokens), float(temperature))
+        api_url, headers, max_tok, temp = get_llm(selected_model, int(max_tokens), float(temperature))
         retriever_obj = retriever(file, int(chunk_size), int(chunk_overlap), embedding_model)
         
         # Get relevant documents
@@ -155,15 +158,28 @@ Question: {query}
 
 Answer:"""
         
-        # Call the model directly
-        response = client.text_generation(
-            prompt,
-            max_new_tokens=max_tok,
-            temperature=temp,
-            return_full_text=False
-        )
+        # Call HuggingFace Inference API directly
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "max_new_tokens": max_tok,
+                "temperature": temp,
+                "return_full_text": False
+            }
+        }
+        
+        response = requests.post(api_url, headers=headers, json=payload)
+        response.raise_for_status()
+        
+        result = response.json()
         
-        return response
+        # Handle different response formats
+        if isinstance(result, list) and len(result) > 0:
+            return result[0].get("generated_text", str(result))
+        elif isinstance(result, dict):
+            return result.get("generated_text", str(result))
+        else:
+            return str(result)
     except Exception as e:
         import traceback
         error_details = traceback.format_exc()