Update app.py

app.py

@@ -1,95 +1,33 @@
-from datetime import datetime
-from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
-from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-from llama_parse import LlamaParse
-from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
-import os
-from dotenv import load_dotenv
-import gradio as gr
-import markdowm as md
-import base64
-
-# Load environment variables
-load_dotenv()
-
-llm_models = [
-    "mistralai/Mixtral-8x7B-Instruct-v0.1",
-    "meta-llama/Meta-Llama-3-8B-Instruct",
-    "mistralai/Mistral-7B-Instruct-v0.2",
-    "tiiuae/falcon-7b-instruct",
-]
-
-embed_models = [
-    "BAAI/bge-small-en-v1.5",  # 33.4M
-    "NeuML/pubmedbert-base-embeddings",
-    "BAAI/llm-embedder", # 109M
-    "BAAI/bge-large-en" # 335M
-]
-
-# Global variable for selected model
-selected_llm_model_name = llm_models[0]  # Default to the first model in the list
-selected_embed_model_name = embed_models[0] # Default to the first model in the list
-vector_index = None
-
-# Initialize the parser
-parser = LlamaParse(api_key=os.getenv("LLAMA_INDEX_API"), result_type='markdown')
-# Define file extractor with various common extensions
-file_extractor = {
-    '.pdf': parser,  # PDF documents
-    '.docx': parser,  # Microsoft Word documents
-    '.doc': parser,  # Older Microsoft Word documents
-    '.txt': parser,  # Plain text files
-    '.csv': parser,  # Comma-separated values files
-    '.xlsx': parser,  # Microsoft Excel files (requires additional processing for tables)
-    '.pptx': parser,  # Microsoft PowerPoint files (for slides)
-    '.html': parser,  # HTML files (web pages)
-
-    # Image files for OCR processing
-    '.jpg': parser,  # JPEG images
-    '.jpeg': parser,  # JPEG images
-    '.png': parser,  # PNG images
-
-
-    # Scanned documents in image formats
-    '.webp': parser,  # WebP images
-    '.svg': parser,  # SVG files (vector format, may contain embedded text)
+# Mapping for display names and actual model names
+llm_display_names = {
+    "tiiuae/falcon-7b-instruct": "HundAI",
+    "mistralai/Mixtral-8x7B-Instruct-v0.1": "Mixtral-8x7B",
+    "meta-llama/Meta-Llama-3-8B-Instruct": "Meta-Llama-3",
+    "mistralai/Mistral-7B-Instruct-v0.2": "Mistral-7B",
 }
 
+# Reverse mapping to retrieve original names
+llm_reverse_mapping = {v: k for k, v in llm_display_names.items()}
 
-# 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)
-        print(f"Parsing done for {file_path}")
-        filename = os.path.basename(file_path)
-        return f"Ready to give response on {filename}"
-    except Exception as e:
-        return f"An error occurred: {e}"
-
-
-# Function to handle the selected model from dropdown
-def set_llm_model(selected_model):
+# Update UI to use display names
+def set_llm_model(display_name):
     global selected_llm_model_name
-    selected_llm_model_name = selected_model  # Update the global variable
-    # print(f"Model selected: {selected_model_name}")
-    # return f"Model set to: {selected_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}")
 
-# Respond function that uses the globally set selected model
+# Respond function remains unchanged
 def respond(message, history):
     try:
         # Initialize the LLM with the selected model
         llm = HuggingFaceInferenceAPI(
-            model_name=selected_llm_model_name,
-            contextWindow=8192,  # Context window size (typically max length of the model)
-            maxTokens=1024,  # Tokens per response generation (512-1024 works well for detailed answers)
-            temperature=0.3,  # Lower temperature for more focused answers (0.2-0.4 for factual info)
-            topP=0.9,  # Top-p sampling to control diversity while retaining quality
-            frequencyPenalty=0.5,  # Slight penalty to avoid repetition
-            presencePenalty=0.5,  # Encourages exploration without digressing too much
+            model_name=selected_llm_model_name,  # Use the backend model name
+            contextWindow=8192,  
+            maxTokens=1024,
+            temperature=0.3,
+            topP=0.9,
+            frequencyPenalty=0.5,
+            presencePenalty=0.5,
             token=os.getenv("TOKEN")
         )
 
@@ -104,10 +42,6 @@ def respond(message, history):
             return "Please upload a file."
         return f"An error occurred: {e}"
 
-def encode_image(image_path):
-    with open(image_path, "rb") as image_file:
-        return base64.b64encode(image_file.read()).decode('utf-8')
-
 # UI Setup
 with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Roboto Mono")]), css='footer {visibility: hidden}') as demo:
     gr.Markdown("")
@@ -121,20 +55,24 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Roboto Mono")]),
             with gr.Row():
                 with gr.Column(scale=1):
                     file_input = gr.File(file_count="single", type='filepath', label="Upload document")
-                    # gr.Markdown("Dont know what to select check out in Intro tab")
                     embed_model_dropdown = gr.Dropdown(embed_models, label="Select Embedding", interactive=True)
                     with gr.Row():
                         btn = gr.Button("Submit", variant='primary')
                         clear = gr.ClearButton()
                     output = gr.Text(label='Vector Index')
-                    llm_model_dropdown = gr.Dropdown(llm_models, label="Select LLM", interactive=True)
+                    
+                    # Use display names for LLM dropdown
+                    llm_model_dropdown = gr.Dropdown(
+                        list(llm_display_names.values()),  # Display names
+                        label="Select LLM",
+                        interactive=True
+                    )
+                    
                 with gr.Column(scale=3):
                     gr.ChatInterface(
                         fn=respond,
                         chatbot=gr.Chatbot(height=500),
-                        theme = "soft",
                         show_progress='full',
-                        # cache_mode='lazy',
                         textbox=gr.Textbox(placeholder="Ask me any questions on the uploaded document!", container=False)
                     )
     
@@ -145,4 +83,4 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Roboto Mono")]),
 
 # Launch the demo with a public link option
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()