Update app.py

app.py

@@ -8,13 +8,10 @@ from transformers import AutoModel, AutoTokenizer
 from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
 from parler_tts import ParlerTTSForConditionalGeneration
 import soundfile as sf
-from llama_index import VectorStoreIndex, SimpleDirectoryReader, LLMPredictor, PromptHelper
-from llama_index.embeddings import GroqEmbedding
-from llama_index.llms import GroqLLM
 from llama_index.agent import ReActAgent
 from llama_index.tools import FunctionTool
+from llama_index.llms import GroqLLM
 from PIL import Image
-from decord import VideoReader, cpu
 from tavily import TavilyClient
 import requests
 from huggingface_hub import hf_hub_download
@@ -85,34 +82,8 @@ def image_generation(query):
     image.save("output.jpg")
     return "output.jpg"
 
-# Document Question Answering Tool
-def doc_question_answering(query, file_path):
-    # Load documents
-    documents = SimpleDirectoryReader(input_files=[file_path]).load_data()
-    
-    # Initialize Groq embedding model
-    embed_model = GroqEmbedding()
-    
-    # Initialize Groq LLM
-    llm_predictor = LLMPredictor(llm=GroqLLM(model_name=MODEL))
-    
-    # Initialize prompt helper
-    prompt_helper = PromptHelper()
-    
-    # Create index
-    index = VectorStoreIndex.from_documents(
-        documents, 
-        embed_model=embed_model,
-        llm_predictor=llm_predictor,
-        prompt_helper=prompt_helper
-    )
-    
-    # Query the index
-    response = index.query(query)
-    return response.response
-
 # Function to handle different input types and choose the right tool
-def handle_input(user_prompt, image=None, video=None, audio=None, doc=None, websearch=False):
+def handle_input(user_prompt, image=None, audio=None, websearch=False):
     if audio:
         if isinstance(audio, str):
             audio = open(audio, "rb") 
@@ -128,14 +99,6 @@ def handle_input(user_prompt, image=None, video=None, audio=None, doc=None, webs
         FunctionTool.from_defaults(fn=image_generation, name="Image Generation"),
     ]
 
-    if doc:
-        tools.append(
-            FunctionTool.from_defaults(
-                fn=lambda query: doc_question_answering(query, doc.name),
-                name="Document Question Answering"
-            )
-        )
-
     llm = GroqLLM(model_name=MODEL)
     agent = ReActAgent.from_tools(tools, llm=llm, verbose=True)
 
@@ -162,7 +125,6 @@ def create_ui():
             with gr.Column(scale=1):
                 image_input = gr.Image(type="filepath", label="Upload an image", elem_id="image-icon")
                 audio_input = gr.Audio(type="filepath", label="Upload audio", elem_id="mic-icon")
-                doc_input = gr.File(type="filepath", label="Upload a document", elem_id="document-icon")
                 voice_only_mode = gr.Checkbox(label="Enable Voice Only Mode", elem_id="voice-only-mode")
                 websearch_mode = gr.Checkbox(label="Enable Web Search", elem_id="websearch-mode")
             with gr.Column(scale=1):
@@ -173,14 +135,14 @@ def create_ui():
 
         submit.click(
             fn=main_interface,
-            inputs=[user_prompt, image_input, audio_input, doc_input, voice_only_mode, websearch_mode],
+            inputs=[user_prompt, image_input, audio_input, voice_only_mode, websearch_mode],
             outputs=[output_label, audio_output]
         )
 
         voice_only_mode.change(
             lambda x: gr.update(visible=not x),
             inputs=voice_only_mode,
-            outputs=[user_prompt, image_input, doc_input, websearch_mode, submit]
+            outputs=[user_prompt, image_input, websearch_mode, submit]
         )
         voice_only_mode.change(
             lambda x: gr.update(visible=x),
@@ -192,13 +154,13 @@ def create_ui():
 
 # Main interface function
 @spaces.GPU()
-def main_interface(user_prompt, image=None, audio=None, doc=None, voice_only=False, websearch=False):
+def main_interface(user_prompt, image=None, audio=None, voice_only=False, websearch=False):
     vqa_model.to(device='cuda', dtype=torch.bfloat16)
     tts_model.to("cuda")
     unet.to("cuda")
     image_pipe.to("cuda")
 
-    response = handle_input(user_prompt, image=image, audio=audio, doc=doc, websearch=websearch)
+    response = handle_input(user_prompt, image=image, audio=audio, websearch=websearch)
 
     if voice_only:
         audio_output = play_voice_output(response)