Update app.py

app.py

@@ -4,23 +4,11 @@ from langchain_community.tools import TavilySearchResults, JinaSearch
 import concurrent.futures
 import json
 import os
-
-UPLOAD_FOLDER = 'uploads'
-if not os.path.exists(UPLOAD_FOLDER):
-    os.makedirs(UPLOAD_FOLDER)
-UPLOAD_FOLDER = 'static'
-if not os.path.exists(UPLOAD_FOLDER):
-    os.makedirs(UPLOAD_FOLDER)
-
 import arxiv 
-import fitz  # PyMuPDF
 from docx import Document
 from PIL import Image
 import io
 import base64
-import mimetypes
-
-
 
 # Set environment variables for Tavily API
 os.environ["TAVILY_API_KEY"] = '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'
@@ -34,55 +22,26 @@ client_3 = Mistral(api_key='cvyu5Rdk2lS026epqL4VB6BMPUcUMSgt')
 def encode_image_bytes(image_bytes):
     return base64.b64encode(image_bytes).decode('utf-8')
 
-# Functions to process various file types
-def process_file(file_path):
-    mime_type, _ = mimetypes.guess_type(file_path)
-    if mime_type == 'application/pdf':
-        return process_pdf(file_path)
-    elif mime_type == 'application/vnd.openxmlformats-officedocument.wordprocessingml.document':
-        return process_docx(file_path)
-    elif mime_type == 'text/plain':
-        return process_txt(file_path)
-    else:
-        print(f"Unsupported file type: {mime_type}")
-        return None, []
-
-def process_pdf(file_path):
-    text = ""
-    images = []
-    pdf_document = fitz.open(file_path)
-    for page_num in range(len(pdf_document)):
-        text += pdf_document[page_num].get_text("text")
-        for _, img in enumerate(pdf_document.get_page_images(page_num, full=True)):
-            xref = img[0]
-            base_image = pdf_document.extract_image(xref)
-            image_bytes = base_image["image"]
-            image_ext = base_image["ext"]
-            base64_image = encode_image_bytes(image_bytes)
-            image_data = f"data:image/{image_ext};base64,{base64_image}"
-            images.append({"type": "image_url", "image_url": image_data})
-    return text, images
-
-def process_docx(file_path):
-    doc = Document(file_path)
-    text = ""
+# Function to decode base64 images
+def decode_base64_image(base64_str):
+    image_data = base64.b64decode(base64_str)
+    return Image.open(io.BytesIO(image_data))
+
+# Process text and images provided by the user
+def process_input(text_input, images_base64):
     images = []
-    for paragraph in doc.paragraphs:
-        text += paragraph.text + "\n"
-    for rel in doc.part.rels.values():
-        if "image" in rel.target_ref:
-            img_data = rel.target_part.blob
-            img = Image.open(io.BytesIO(img_data))
-            buffered = io.BytesIO()
-            img.save(buffered, format="JPEG")
-            image_base64 = encode_image_bytes(buffered.getvalue())
-            images.append({"type": "image_url", "image_url": f"data:image/jpeg;base64,{image_base64}"})
-    return text, images
-
-def process_txt(file_path):
-    with open(file_path, "r", encoding="utf-8") as file:
-        text = file.read()
-    return text, []
+    if images_base64:
+        for img_data in images_base64:
+            try:
+                img = decode_base64_image(img_data)
+                buffered = io.BytesIO()
+                img.save(buffered, format="JPEG")
+                image_base64 = encode_image_bytes(buffered.getvalue())
+                images.append({"type": "image_url", "image_url": f"data:image/jpeg;base64,{image_base64}"})
+            except Exception as e:
+                print(f"Error decoding image: {e}")
+
+    return text_input, images
 
 # Search setup function
 def setup_search(question):
@@ -204,11 +163,8 @@ def ask_question_to_mistral(text, question, images=[]):
     return response.choices[0].message.content
 
 # Gradio interface
-def gradio_interface(file, task, question, compression_percentage):
-    if file:
-        text, images = process_file(file.name)
-    else:
-        text, images = "", []
+def gradio_interface(text_input, images_base64, task, question, compression_percentage):
+    text, images = process_input(text_input, images_base64)
 
     topics, articles_json = init(text, images)
 
@@ -225,7 +181,8 @@ def gradio_interface(file, task, question, compression_percentage):
 with gr.Blocks() as demo:
     gr.Markdown("## Text Analysis: Summarization or Question Answering")
     with gr.Row():
-        file_input = gr.File(label="Upload File")
+        text_input = gr.Textbox(label="Input Text")
+        images_base64 = gr.Textbox(label="Base64 Images (comma-separated, if any)", placeholder="data:image/jpeg;base64,...", lines=2)
         task_choice = gr.Radio(["Summarization", "Question Answering"], label="Select Task")
         question_input = gr.Textbox(label="Question (for Question Answering)", visible=False)
         compression_input = gr.Slider(label="Compression Percentage (for Summarization)", minimum=10, maximum=90, value=30, visible=False)
@@ -238,6 +195,6 @@ with gr.Blocks() as demo:
         result_output = gr.JSON(label="Results")
 
     submit_button = gr.Button("Submit")
-    submit_button.click(gradio_interface, [file_input, task_choice, question_input, compression_input], result_output)
+    submit_button.click(gradio_interface, [text_input, images_base64, task_choice, question_input, compression_input], result_output)
 
 demo.launch()