Update app.py

app.py

@@ -19,6 +19,13 @@ import inspect
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
 llama_cloud_api_key = os.environ.get("LLAMA_CLOUD_API_KEY")
 
+MODELS = [
+    "meta-llama/Meta-Llama-3.1-70B-Instruct",
+    "meta-llama/Meta-Llama-3-70B-Instruct",
+    "mistralai/Mistral-7B-Instruct-v0.3",
+    "mistralai/Mixtral-8x7B-Instruct-v0.1"
+]
+
 # Initialize LlamaParse
 llama_parser = LlamaParse(
     api_key=llama_cloud_api_key,
@@ -71,9 +78,9 @@ def update_vectors(files, parser):
     
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
 
-def generate_chunked_response(prompt, max_tokens=1000, max_chunks=5, temperature=0.7):
+def generate_chunked_response(prompt, model, max_tokens=1000, max_chunks=5, temperature=0.7):
     client = InferenceClient(
-        "meta-llama/Meta-Llama-3.1-70B-Instruct",
+        model,
         token=huggingface_token,
     )
     
@@ -122,7 +129,7 @@ class CitingSources(BaseModel):
         description="List of sources to cite. Should be an URL of the source."
     )
 
-def get_response_from_pdf(query, temperature=0.7):
+def get_response_from_pdf(query, model, temperature=0.7):
     embed = get_embeddings()
     if os.path.exists("faiss_database"):
         database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
@@ -146,7 +153,7 @@ Do not include a list of sources in your response. [/INST]"""
 
     return clean_text
 
-def get_response_with_search(query, temperature=0.7):
+def get_response_with_search(query, model, temperature=0.7):
     search_results = duckduckgo_search(query)
     context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
                         for result in search_results if 'body' in result)
@@ -169,12 +176,12 @@ After writing the document, please provide a list of sources used in your respon
     
     return main_content, sources
 
-def chatbot_interface(message, history, use_web_search, temperature):
+def chatbot_interface(message, history, use_web_search, model, temperature):
     if use_web_search:
-        main_content, sources = get_response_with_search(message, temperature)
+        main_content, sources = get_response_with_search(message, model, temperature)
         formatted_response = f"{main_content}\n\nSources:\n{sources}"
     else:
-        response = get_response_from_pdf(message, temperature)
+        response = get_response_from_pdf(message, model, temperature)
         formatted_response = response
 
     history.append((message, formatted_response))
@@ -182,7 +189,7 @@ def chatbot_interface(message, history, use_web_search, temperature):
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# AI-powered Web Search and PDF Chat Assistant (Using Meta-Llama-3.1-8B-Instruct)")
+    gr.Markdown("# AI-powered Web Search and PDF Chat Assistant")
     
     with gr.Row():
         file_input = gr.Files(label="Upload your PDF documents", file_types=[".pdf"])
@@ -197,7 +204,8 @@ with gr.Blocks() as demo:
     use_web_search = gr.Checkbox(label="Use Web Search", value=False)
     
     with gr.Row():
-        temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.3, step=0.1, label="Temperature")
+        model_dropdown = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[2])
+        temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
     
     submit = gr.Button("Submit")
 
@@ -212,10 +220,10 @@ with gr.Blocks() as demo:
     )
 
     submit.click(chatbot_interface, 
-                 inputs=[msg, chatbot, use_web_search, temperature_slider], 
+                 inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider], 
                  outputs=[chatbot])
     msg.submit(chatbot_interface, 
-               inputs=[msg, chatbot, use_web_search, temperature_slider], 
+               inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider], 
                outputs=[chatbot])
 
     gr.Markdown(