Removed ChatHuggingface with custom wrapper to wrap InferenceClient

app.py

@@ -12,11 +12,42 @@ from langchain.text_splitter import CharacterTextSplitter
 from huggingface_hub import InferenceClient
 import logging
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# logging.basicConfig(level=logging.INFO)
+# logger = logging.getLogger(__name__)
 
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
+class HuggingFaceInterferenceClientRunnable(Runnable):
+    def __init__(self, client, max_tokens=512, temperature=0.7, top_p=0.95):
+        self.client = client
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+        self.top_p = top_p
+
+    def invoke(self, input, config=None):
+        prompt = input.to_messages()[0].content
+        messages = [{"role": "user", "content": prompt}]
+
+        response = ""
+        for part in self.client.chat_completion(
+            messages,
+            max_tokens=self.max_tokens,
+            stream=True,
+            temperature=self.temperature,
+            top_p=self.top_p
+        ):
+            token = part.choices[0].delta.content
+            if token:
+                response += token
+
+        return response
+
+    def update_params(self, max_tokens, temperature, top_p):
+        self.max_tokens = max_tokens
+        self.temperature=temperature
+        self.top_p=top_p
+
+
 def extract_pdf_text(url: str) -> str:
     response = requests.get(url)
     pdf_file = BytesIO(response.content)
@@ -39,11 +70,7 @@ vectorstore = Chroma.from_documents(
 )
 retriever = vectorstore.as_retriever()
 
-llm = ChatHuggingFace(
-    huggingfacehub_api_token=None,
-    model_id="HuggingFaceH4/zephyr-7b-beta",
-    interference_client=client,
-)
+llm = HuggingFaceInterferenceClientRunnable(client)
 
 # Before RAG chain
 before_rag_template = "What is {topic}"
@@ -75,9 +102,8 @@ after_rag_chain = (
 )
 
 def process_query(role, system_message, max_tokens, temperature, top_p):
-    client.max_tokens = max_tokens
-    client.temperature = temperature
-    client.top_p = top_p
+
+    llm.update_params(max_tokens, temperature, top_p)
 
     # Before RAG
     before_rag_result = before_rag_chain.invoke({"topic": "Hugging Face"})