Update app.py

app.py

@@ -6,13 +6,17 @@ 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()
 
-models = [
+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",
     # "NousResearch/Yarn-Mistral-7b-64k",  ## 14GB>10GB
     # "impira/layoutlm-document-qa",  ## ERR
     # "Qwen/Qwen1.5-7B",      ## 15GB
@@ -20,7 +24,6 @@ models = [
     # "google/gemma-2-2b-jpn-it",   ## high response time
     # "impira/layoutlm-invoices",   ## bad req
     # "google/pix2struct-docvqa-large",  ## bad req
-    "mistralai/Mistral-7B-Instruct-v0.2",
     # "google/gemma-7b-it", ## 17GB > 10GB
     # "google/gemma-2b-it",  ## high response time
     # "HuggingFaceH4/zephyr-7b-beta",   ## high response time
@@ -28,14 +31,23 @@ models = [
     # "microsoft/phi-2",    ## high response time
     # "TinyLlama/TinyLlama-1.1B-Chat-v1.0",     ## high response time
     # "mosaicml/mpt-7b-instruct",     ## 13GB>10GB
-    "tiiuae/falcon-7b-instruct",
     # "google/flan-t5-xxl" ## high respons time
     # "NousResearch/Yarn-Mistral-7b-128k",  ## 14GB>10GB
     # "Qwen/Qwen2.5-7B-Instruct",     ## 15GB>10GB
 ]
 
+embed_models = [
+    "BAAI/bge-small-en-v1.5",  # 33.4M
+    "NeuML/pubmedbert-base-embeddings",
+    "sentence-transformers/all-mpnet-base-v2" # 109M
+    "BAAI/llm-embedder", # 109M
+    "BAAI/bge-large-en" # 335
+]
+
 # Global variable for selected model
-selected_model_name = models[0]  # Default to the first model in the list
+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')
@@ -67,24 +79,12 @@ file_extractor = {
     '.svg': parser,  # SVG files (vector format, may contain embedded text)
 }
 
-# Embedding model and index initialization (to be populated by uploaded files)
-embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")   ## Works good and fast
-# embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-large-en")    ## works good
-# embed_model2 = HuggingFaceEmbedding(model_name="NeuML/pubmedbert-base-embeddings") ## works good
-
-# sentence-transformers/distilbert-base-nli-mean-tokens
-# BAAI/bge-large-en
-# embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-MiniLM-L6-v2")
-
-# Global variable to store documents loaded from user-uploaded files
-vector_index = None
-
-
 # File processing function
-def load_files(file_path: str):
+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)
@@ -94,9 +94,9 @@ def load_files(file_path: str):
 
 
 # Function to handle the selected model from dropdown
-def set_model(selected_model):
-    global selected_model_name
-    selected_model_name = selected_model  # Update the global variable
+def set_llm_model(selected_model):
+    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}"
 
@@ -106,51 +106,68 @@ def respond(message, history):
     try:
         # Initialize the LLM with the selected model
         llm = HuggingFaceInferenceAPI(
-            model_name=selected_model_name,
-            contextWindow = 4096,
-            maxTokens = 4096,
-            temperature=0.7,
-            topP=0.95,
-            # token=os.getenv("TOKEN")
+            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
+            token=os.getenv("TOKEN")
         )
 
-        # Check selected model
-        # print(f"Using model: {selected_model_name}")
-
         # 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_model_name}:: {message} --> {str(bot_message)}\n")
-        return f"{selected_model_name}:\n{str(bot_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)}"
     except Exception as e:
         if str(e) == "'NoneType' object has no attribute 'as_query_engine'":
             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')
+
+# Encode the images
+github_logo_encoded = encode_image("Images/github-logo.png")
+linkedin_logo_encoded = encode_image("Images/linkedin-logo.png")
+website_logo_encoded = encode_image("Images/ai-logo.png")
 
 # UI Setup
-with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column(scale=1):
-            file_input = gr.File(file_count="single", type='filepath', label="Step-1: Upload document")
-            with gr.Row():
-                clear = gr.ClearButton()
-                btn = gr.Button("Submit", variant='primary')
-            output = gr.Text(label='Vector Index')
-            model_dropdown = gr.Dropdown(models, label="Step-2: Select Model", interactive=True)
-
-        with gr.Column(scale=3):
-            gr.ChatInterface(
-                fn=respond,
-                chatbot=gr.Chatbot(height=500),
-                textbox=gr.Textbox(placeholder="Step-3: Ask me questions on the uploaded document!", container=False, scale=7)
-            )
+with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Roboto Mono")]), css='footer {visibility: hidden}') as demo:
+    gr.Markdown("# DocBot📄🤖")
+    with gr.Tabs():
+        with gr.TabItem("Intro"):
+            gr.Markdown(md.description)
 
+        with gr.TabItem("DocBot"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    file_input = gr.File(file_count="single", type='filepath', label="Step-1: Upload document")
+                    gr.Markdown("Dont know what to select check out in Intro tab")
+                    embed_model_dropdown = gr.Dropdown(embed_models, label="Step-2: Select Embedding", interactive=True)
+                    with gr.Row():
+                        clear = gr.ClearButton()
+                        btn = gr.Button("Submit", variant='primary')
+                    output = gr.Text(label='Vector Index')
+                    llm_model_dropdown = gr.Dropdown(llm_models, label="Step-3: 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="Step-4: Ask me questions on the uploaded document!", container=False)
+                    )
+    gr.HTML(md.footer.format(github_logo_encoded, linkedin_logo_encoded, website_logo_encoded))
     # Set up Gradio interactions
-    model_dropdown.change(fn=set_model, inputs=model_dropdown)
-    btn.click(fn=load_files, inputs=file_input, outputs=output)
-    clear.click(lambda: [None] * 2, outputs=[file_input, output])
+    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__":