Update app.py

app.py

@@ -7,14 +7,20 @@ import numpy as np
 from sentence_transformers import SentenceTransformer
 import gradio as gr
 
-# Configure Hugging Face API
-huggingface_api_url = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-70B-Instruct"
+# 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"
+}
+
 huggingface_api_key = os.getenv("HUGGINGFACE_API_KEY")
 headers = {"Authorization": f"Bearer {huggingface_api_key}"}
 
 # Function to query Hugging Face model
-def query_huggingface(payload):
-    response = requests.post(huggingface_api_url, headers=headers, json=payload)
+def query_huggingface(model_name, payload):
+    response = requests.post(api_urls[model_name], headers=headers, json=payload)
     return response.json()
 
 # Function to extract text from PDF
@@ -83,12 +89,12 @@ init_db()
 model = SentenceTransformer('all-MiniLM-L6-v2')
 faiss_index, context_list = update_faiss_index()
 
-# Gradio interface
-def chatbot(question):
+# Gradio interface for chatbot
+def chatbot(model_name, question):
     relevant_contexts = retrieve_relevant_context(faiss_index, context_list, question)
     user_input = f"question: {question} context: {' '.join(relevant_contexts)}"
-    response = query_huggingface({"inputs": user_input})
-    response_text = response.get("generated_text", "Sorry, I couldn't generate a response.")
+    response = query_huggingface(model_name, {"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
 
 # File upload function
@@ -100,8 +106,13 @@ def upload_pdf(file):
     return "PDF content added to context."
 
 # Gradio interface
-iface = gr.Interface(fn=chatbot, inputs="text", outputs="text", title="Storage Warehouse Customer Service Chatbot")
+iface = gr.Interface(
+    fn=chatbot,
+    inputs=[gr.inputs.Dropdown(["Meta-Llama-3-70B-Instruct", "Meta-Llama-3-8B-Instruct", "Gemma-2-27B-IT", "Gemma-2-27B"]), "text"],
+    outputs="text",
+    title="Storage Warehouse Customer Service Chatbot"
+)
 file_upload = gr.Interface(fn=upload_pdf, inputs="file", outputs="text", title="Upload PDF for Context")
 
 app = gr.TabbedInterface([iface, file_upload], ["Chatbot", "Upload PDF"])
-app.launch()
+app.launch(share=True)