Update app.py

app.py

@@ -12,37 +12,64 @@ import base64
 # Load environment variables
 load_dotenv()
 
-# Mapping for display names and actual model names
-llm_display_names = {
-    "tiiuae/falcon-7b-instruct": "HundAI",
+llm_models = {
     "mistralai/Mixtral-8x7B-Instruct-v0.1": "Mixtral-8x7B",
-    "meta-llama/Meta-Llama-3-8B-Instruct": "Meta-Llama-3",
+    "meta-llama/Meta-Llama-3-8B-Instruct": "Meta-Llama-8B",
     "mistralai/Mistral-7B-Instruct-v0.2": "Mistral-7B",
+    "tiiuae/falcon-7b-instruct": "HundAI",  # Model renamed for UI display
 }
+
 embed_models = [
     "BAAI/bge-small-en-v1.5",  # 33.4M
     "NeuML/pubmedbert-base-embeddings",
-    "BAAI/llm-embedder",  # 109M
-    "BAAI/bge-large-en"  # 335M
+    "BAAI/llm-embedder", # 109M
+    "BAAI/bge-large-en" # 335M
 ]
 
-# Reverse mapping to retrieve original names
-llm_reverse_mapping = {v: k for k, v in llm_display_names.items()}
+# Global variable for selected model
+selected_llm_model_name = list(llm_models.keys())[0]  # Default to the first model in the dictionary
+vector_index = None
+
+# Initialize the parser
+parser = LlamaParse(api_key=os.getenv("LLAMA_INDEX_API"), result_type='markdown')
+file_extractor = {
+    '.pdf': parser,
+    '.docx': parser,
+    '.txt': parser,
+    '.csv': parser,
+    '.xlsx': parser,
+    '.pptx': parser,
+    '.html': parser,
+    '.jpg': parser,
+    '.jpeg': parser,
+    '.png': parser,
+    '.webp': parser,
+    '.svg': parser,
+}
+
+# File processing function
+def load_files(file_path: str, embed_model_name: str):
+    try:
+        global vector_index
+        document = SimpleDirectoryReader(input_files=[file_path], file_extractor=file_extractor).load_data()
+        embed_model = HuggingFaceEmbedding(model_name=embed_model_name)
+        vector_index = VectorStoreIndex.from_documents(document, embed_model=embed_model)
+        filename = os.path.basename(file_path)
+        return f"Ready to give response on {filename}"
+    except Exception as e:
+        return f"An error occurred: {e}"
 
-# Update UI to use display names
-def set_llm_model(display_name):
+# Function to handle the selected model from dropdown
+def set_llm_model(selected_model):
     global selected_llm_model_name
-    # Retrieve the original model name using the reverse mapping
-    selected_llm_model_name = llm_reverse_mapping.get(display_name, display_name)
-    print(f"Model selected: {selected_llm_model_name}")
+    selected_llm_model_name = next(key for key, value in llm_models.items() if value == selected_model)
 
-# Respond function remains unchanged
+# Respond function
 def respond(message, history):
     try:
-        # Initialize the LLM with the selected model
         llm = HuggingFaceInferenceAPI(
-            model_name=selected_llm_model_name,  # Use the backend model name
-            contextWindow=8192,  
+            model_name=selected_llm_model_name,
+            contextWindow=8192,
             maxTokens=1024,
             temperature=0.3,
             topP=0.9,
@@ -50,13 +77,9 @@ def respond(message, history):
             presencePenalty=0.5,
             token=os.getenv("TOKEN")
         )
-
-        # Set up the query engine with the selected LLM
         query_engine = vector_index.as_query_engine(llm=llm)
         bot_message = query_engine.query(message)
-
-        print(f"\n{datetime.now()}:{selected_llm_model_name}:: {message} --> {str(bot_message)}\n")
-        return f"{selected_llm_model_name}:\n{str(bot_message)}"
+        return f"{llm_models[selected_llm_model_name]}:\n{str(bot_message)}"
     except Exception as e:
         if str(e) == "'NoneType' object has no attribute 'as_query_engine'":
             return "Please upload a file."
@@ -80,27 +103,18 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Roboto Mono")]),
                         btn = gr.Button("Submit", variant='primary')
                         clear = gr.ClearButton()
                     output = gr.Text(label='Vector Index')
-                    
-                    # Use display names for LLM dropdown
-                    llm_model_dropdown = gr.Dropdown(
-                        list(llm_display_names.values()),  # Display names
-                        label="Select LLM",
-                        interactive=True
-                    )
-                    
+                    llm_model_dropdown = gr.Dropdown(list(llm_models.values()), label="Select LLM", interactive=True)
                 with gr.Column(scale=3):
                     gr.ChatInterface(
                         fn=respond,
                         chatbot=gr.Chatbot(height=500),
-                        show_progress='full',
+                        theme="soft",
                         textbox=gr.Textbox(placeholder="Ask me any questions on the uploaded document!", container=False)
                     )
     
-    # Set up Gradio interactions
     llm_model_dropdown.change(fn=set_llm_model, inputs=llm_model_dropdown)
     btn.click(fn=load_files, inputs=[file_input, embed_model_dropdown], outputs=output)
     clear.click(lambda: [None] * 3, outputs=[file_input, embed_model_dropdown, output])
 
-# Launch the demo with a public link option
 if __name__ == "__main__":
     demo.launch()