Update app.py

app.py

@@ -1,11 +1,12 @@
 import os
-from langchain_huggingface import HuggingFaceEndpoint, HuggingFaceEmbeddings
+from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 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
@@ -54,16 +55,9 @@ def get_huggingface_token():
 def get_llm(model_id: str = MODEL_OPTIONS[0], max_tokens: int = 256, temperature: float = 0.8):
     token = get_huggingface_token()
     
-    # Use simpler initialization without specifying task
-    # Let HuggingFace auto-detect the best configuration
-    llm = HuggingFaceEndpoint(
-        repo_id=model_id,
-        max_new_tokens=max_tokens,
-        temperature=temperature,
-        huggingfacehub_api_token=token,
-        timeout=120,  # Increase timeout for large models
-    )
-    return llm
+    # Use InferenceClient directly for better reliability
+    client = InferenceClient(model=model_id, token=token)
+    return client, max_tokens, temperature
 
 
 # ---------------------------
@@ -146,31 +140,29 @@ def retriever_qa(file, query, model_choice, max_tokens, temperature, embedding_m
     
     try:
         selected_model = model_choice or MODEL_OPTIONS[0]
-        llm = get_llm(selected_model, int(max_tokens), float(temperature))
+        client, max_tok, temp = get_llm(selected_model, int(max_tokens), float(temperature))
         retriever_obj = retriever(file, int(chunk_size), int(chunk_overlap), embedding_model)
         
-        # Create a simple RAG chain
-        template = """Answer the question based only on the following context:
+        # Get relevant documents
+        docs = retriever_obj.get_relevant_documents(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:
 {context}
 
-Question: {question}
+Question: {query}
 
 Answer:"""
         
-        prompt = PromptTemplate.from_template(template)
-        
-        def format_docs(docs):
-            return "\n\n".join(doc.page_content for doc in docs)
-        
-        # Build the chain
-        rag_chain = (
-            {"context": retriever_obj | format_docs, "question": RunnablePassthrough()}
-            | prompt
-            | llm
-            | StrOutputParser()
+        # Call the model directly
+        response = client.text_generation(
+            prompt,
+            max_new_tokens=max_tok,
+            temperature=temp,
+            return_full_text=False
         )
         
-        response = rag_chain.invoke(query)
         return response
     except Exception as e:
         import traceback