Update app.py

app.py

@@ -1,5 +1,5 @@
 import os
-import requests
+from huggingface_hub import InferenceClient
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
@@ -55,16 +55,12 @@ def get_huggingface_token():
 # ---------------------------
 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.
+    Returns InferenceClient for HuggingFace models.
     """
     token = get_huggingface_token()
-    api_url = f"https://api-inference.huggingface.co/models/{model_id}"
-    headers = {
-        "Authorization": f"Bearer {token}",
-        "Content-Type": "application/json"
-    }
+    client = InferenceClient(token=token)
     
-    return api_url, headers, max_tokens, temperature
+    return client, model_id, max_tokens, temperature
 
 
 # ---------------------------
@@ -147,43 +143,39 @@ def retriever_qa(file, query, model_choice, max_tokens, temperature, embedding_m
     
     try:
         selected_model = model_choice or MODEL_OPTIONS[0]
-        api_url, headers, max_tok, temp = get_llm(selected_model, int(max_tokens), float(temperature))
+        client, model_id, 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
         docs = retriever_obj.invoke(query)
         context = "\n\n".join(doc.page_content for doc in docs)
         
-        # Create prompt
-        prompt = f"""Answer the question based only on the following context:
+        # Create messages for chat completion
+        messages = [
+            {
+                "role": "system",
+                "content": "You are a helpful assistant that answers questions based only on the provided context."
+            },
+            {
+                "role": "user",
+                "content": f"""Context:
 {context}
 
 Question: {query}
 
-Answer:"""
-        
-        # Call HuggingFace Inference API directly
-        payload = {
-            "inputs": prompt,
-            "parameters": {
-                "max_new_tokens": max_tok,
-                "temperature": temp,
-                "return_full_text": False
+Please answer the question based only on the context provided above."""
             }
-        }
-        
-        response = requests.post(api_url, headers=headers, json=payload)
-        response.raise_for_status()
+        ]
         
-        result = response.json()
+        # Call chat completion API
+        response = client.chat_completion(
+            messages=messages,
+            model=model_id,
+            max_tokens=max_tok,
+            temperature=temp
+        )
         
-        # 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)
+        return response.choices[0].message.content
     except Exception as e:
         import traceback
         error_details = traceback.format_exc()