Update app.py

app.py

@@ -1,3 +1,4 @@
+from huggingface_hub import InferenceClient
 import os
 import sqlite3
 import requests
@@ -7,20 +8,14 @@ import numpy as np
 from sentence_transformers import SentenceTransformer
 import gradio as gr
 
-# Configure Hugging Face API URLs and headers
-api_urls = {
-    "Meta-Llama-3-70B-Instruct": "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-70B-Instruct",
-    "Meta-Llama-3-8B-Instruct": "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B-Instruct",
-    "Gemma-2-27B-IT": "https://api-inference.huggingface.co/models/google/gemma-2-27b-it",
-    "Gemma-2-27B": "https://api-inference.huggingface.co/models/google/gemma-2-27b"
-}
-
+# Configure Hugging Face API URL and headers
+model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
 huggingface_api_key = os.getenv("HUGGINGFACE_API_KEY")
 headers = {"Authorization": f"Bearer {huggingface_api_key}"}
 
 # Function to query Hugging Face model
-def query_huggingface(model_name, payload):
-    response = requests.post(api_urls[model_name], headers=headers, json=payload)
+def query_huggingface(payload):
+    response = requests.post(f"https://api-inference.huggingface.co/models/{model_name}", headers=headers, json=payload)
     return response.json()
 
 # Function to extract text from PDF
@@ -90,10 +85,10 @@ model = SentenceTransformer('all-MiniLM-L6-v2')
 faiss_index, context_list = update_faiss_index()
 
 # Gradio interface for chatbot
-def chatbot(model_name, question):
+def chatbot(question):
     relevant_contexts = retrieve_relevant_context(faiss_index, context_list, question)
     user_input = f"question: {question} context: {' '.join(relevant_contexts)}"
-    response = query_huggingface(model_name, {"inputs": user_input})
+    response = query_huggingface({"inputs": user_input})
     response_text = response[0].get("generated_text", "Sorry, I couldn't generate a response.") if isinstance(response, list) else response.get("generated_text", "Sorry, I couldn't generate a response.")
     return response_text
 
@@ -108,7 +103,7 @@ def upload_pdf(file):
 # Gradio interface
 iface = gr.Interface(
     fn=chatbot,
-    inputs=[gr.Dropdown(["Meta-Llama-3-70B-Instruct", "Meta-Llama-3-8B-Instruct", "Gemma-2-27B-IT", "Gemma-2-27B"]), gr.Textbox()],
+    inputs=gr.Textbox(),
     outputs=gr.Textbox(),
     title="Storage Warehouse Customer Service Chatbot"
 )