Update app.py

app.py

@@ -3,6 +3,8 @@ import os
 import sys
 
 from fastapi import FastAPI, File, UploadFile
+import gradio as gr
+import requests
 from typing import List
 import torch
 from pdf2image import convert_from_path
@@ -80,6 +82,47 @@ async def search(query: str, k: int):
 
     return {"results": results}
 
+def index_gradio(file, ds):
+    """Upload PDFs and get embeddings."""
+    url = "http://localhost:8082/index"
+    files = [("files", (f.name, f.file)) for f in file]
+    response = requests.post(url, files=files)
+    result = response.json()
+    return result['message'], ds, []
+
+def search_gradio(query: str, ds, images, k):
+    """Send a search query and get results."""
+    url = "http://localhost:8082/search"
+    payload = {'query': query, 'k': k}
+    response = requests.post(url, json=payload)
+    result = response.json()
+    results = result['results']
+    return results
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# ColPali: Efficient Document Retrieval with Vision Language Models 📚")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            gr.Markdown("## 1️⃣ Upload PDFs")
+            file = gr.File(file_types=["pdf"], file_count="multiple", label="Upload PDFs")
+
+            convert_button = gr.Button("🔄 Index documents")
+            message = gr.Textbox("Files not yet uploaded", label="Status")
+            embeds = gr.State(value=[])
+            imgs = gr.State(value=[])
+
+        with gr.Column(scale=3):
+            gr.Markdown("## 2️⃣ Search")
+            query = gr.Textbox(placeholder="Enter your query here", label="Query")
+            k = gr.Slider(minimum=1, maximum=10, step=1, label="Number of results", value=1)
+
+    # Define the actions
+    search_button = gr.Button("🔍 Search", variant="primary")
+    output_gallery = gr.Gallery(label="Retrieved Documents", height=600, show_label=True)
+
+    convert_button.click(index_gradio, inputs=[file, embeds], outputs=[message, embeds, imgs])
+    search_button.click(search_gradio, inputs=[query, embeds, imgs, k], outputs=[output_gallery])
+
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8082, reload=True)
+    demo.queue(max_size=10).launch(debug=True)