Update app.py

app.py

@@ -6,132 +6,73 @@ 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",
-]
+# Initialize the LLM and parser
+llm = HuggingFaceInferenceAPI(
+    model_name="meta-llama/Meta-Llama-3-8B-Instruct",
+    token=os.getenv("TOKEN")
+)
 
-embed_models = [
-    "BAAI/bge-small-en-v1.5",
-    "NeuML/pubmedbert-base-embeddings",
-    "BAAI/llm-embedder",
-    "BAAI/bge-large-en"
-]
+parser = LlamaParse(api_key=os.getenv("LLAMA_INDEX_API"), result_type='markdown')
+file_extractor = {'.pdf': parser, '.docx': parser, '.doc': parser}
 
-# Global state
-selected_llm_model_name = llm_models[0]
-selected_embed_model_name = embed_models[0]
-vector_index = None
+# Embedding model and index initialization (to be populated by uploaded files)
+embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
 
-# Parser setup
-parser = LlamaParse(api_key=os.getenv("LLAMA_INDEX_API"), result_type='markdown')
-file_extractor = {ext: parser for ext in ['.pdf', '.docx', '.doc', '.txt', '.csv', '.xlsx', '.pptx', '.html', '.jpg', '.jpeg', '.png', '.webp', '.svg']}
+# Global variable to store documents loaded from user-uploaded files
+vector_index = None
 
-def load_files(file_path: str, embed_model_name: str):
-    global vector_index
+# File processing function
+def load_files(file_path: 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 {file_path}")
         filename = os.path.basename(file_path)
-        return f"Ready to give response on {filename}"
+        return f"Ready to give response on give {filename}"
     except Exception as e:
-        return f"An error occurred: {e}"
-
-def set_llm_model(selected_model):
-    global selected_llm_model_name
-    selected_llm_model_name = selected_model
-    return f"Model set to: {selected_model}"
+        return f"An error occurred {e}"
 
 def respond(message, history):
     try:
-        llm = HuggingFaceInferenceAPI(
-            model_name=selected_llm_model_name,
-            contextWindow=8192,
-            maxTokens=1024,
-            temperature=0.3,
-            topP=0.9,
-            frequencyPenalty=0.5,
-            presencePenalty=0.5,
-            token=os.getenv("TOKEN")
-        )
-        if vector_index is not None:
-            query_engine = vector_index.as_query_engine(llm=llm)
-            bot_message = str(query_engine.query(message))
-            history.append((message, bot_message))
-            print(f"\n{datetime.now()}:{selected_llm_model_name}:: {message} --> {bot_message}\n")
-            return bot_message, history
-        else:
-            return "Please upload a file first.", history
+        query_engine = vector_index.as_query_engine(llm=llm)
+        bot_message = query_engine.query(message)
+        # output = ""
+        # for chr in bot_message:
+        #     output += chr
+        #     yield output
+        print(f"{datetime.now()}::message=>{str(bot_message)}")
+        return str(bot_message)
     except Exception as e:
-        return f"An error occurred: {e}", history
-
-def encode_image(image_path):
-    with open(image_path, "rb") as image_file:
-        return base64.b64encode(image_file.read()).decode('utf-8')
-
-# Encoded logos
-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")
-
-# Markdown placeholders
-description = "### Welcome to **DocBot** - Ask Questions Based on Your Uploaded Documents"
-guide = "> Step 1: Upload\n> Step 2: Select Embedding\n> Step 3: Select LLM\n> Step 4: Ask Questions"
-footer = """
-<center>
-    <a href="https://github.com" target="_blank"><img src="data:image/png;base64,{}" height="30"/></a>&nbsp;
-    <a href="https://linkedin.com" target="_blank"><img src="data:image/png;base64,{}" height="30"/></a>&nbsp;
-    <a href="https://yourwebsite.com" target="_blank"><img src="data:image/png;base64,{}" height="30"/></a>
-</center>
-""".format(github_logo_encoded, linkedin_logo_encoded, website_logo_encoded)
-
-# Gradio UI
-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(description)
-
-        with gr.TabItem("DocBot"):
-            with gr.Accordion("=== IMPORTANT: READ ME FIRST ===", open=False):
-                gr.Markdown(guide)
-
-            with gr.Row():
-                with gr.Column(scale=1):
-                    file_input = gr.File(file_count="single", type='filepath', label="Step-1: Upload document")
-                    embed_model_dropdown = gr.Dropdown(embed_models, label="Step-2: 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="Step-3: Select LLM", interactive=True)
-                    model_selected_output = gr.Text(label="Model selected")
-
-                with gr.Column(scale=3):
-                    chatbot_ui = gr.Chatbot(height=500)
-                    message = gr.Textbox(placeholder="Step-4: Ask me questions on the uploaded document!", container=False)
-                    submit_btn = gr.Button("Send")
-
-            # Bind logic
-            llm_model_dropdown.change(fn=set_llm_model, inputs=llm_model_dropdown, outputs=model_selected_output)
-            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])
-
-            # Chat logic
-            state = gr.State([])
-            submit_btn.click(fn=respond, inputs=[message, state], outputs=[chatbot_ui, state])
-            message.submit(fn=respond, inputs=[message, state], outputs=[chatbot_ui, state])
-
-    gr.HTML(footer)
-
+        if e == "'NoneType' object has no attribute 'as_query_engine'":
+            return "upload file"
+        return f"an error occurred {e}"
+
+# UI Setup
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column(scale=1):
+            file_input = gr.File(file_count="single", type='filepath')
+            with gr.Column():
+                clear = gr.ClearButton()
+                btn = gr.Button("Submit", variant='primary')
+            output = gr.Text(label='Vector Index')
+        with gr.Column(scale=2):
+             gr.ChatInterface(fn=respond,
+                              chatbot=gr.Chatbot(height=500),
+                              textbox=gr.Textbox(placeholder="Ask me a yes or no question", container=False, scale=7),
+                              examples=["summarize the document"]
+                             )
+
+    # Action on button click to process file and load into index
+    btn.click(fn=load_files, inputs=file_input, outputs=output)
+    clear.click(lambda: [None]*2, outputs=[file_input, output])
+
+
+# Launch the demo with public link option
 if __name__ == "__main__":
     demo.launch(share=True)