chore: Update requirements.txt with new dependencies

app.py

@@ -1,15 +1,30 @@
+import asyncio
+import html
 import os
-import torch
+from io import BytesIO
+
+import aiohttp
+import dotenv
 import gradio as gr
 import requests
+import torch
+from colpali_engine.models import ColQwen2, ColQwen2Processor
 from PIL import Image
-from io import BytesIO
 from qdrant_client import QdrantClient
-from colpali_engine.models import ColQwen2, ColQwen2Processor
 
+dotenv.load_dotenv()
+
+if torch.cuda.is_available():
+    device = "cuda:0"
+elif torch.backends.mps.is_available():
+    device = "mps"
+else:
+    device = "cpu"
+
+
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 # Initialize ColPali model and processor
 model_name = "vidore/colqwen2-v0.1"
-device = "cuda:0" if torch.cuda.is_available() else "cpu"  # You can change this to "mps" for Apple Silicon if needed
 colpali_model = ColQwen2.from_pretrained(
     model_name,
     torch_dtype=torch.bfloat16,
@@ -21,15 +36,22 @@ colpali_processor = ColQwen2Processor.from_pretrained(
 
 # Initialize Qdrant client
 QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
-qdrant_client = QdrantClient(url="https://davanstrien-qdrant-test.hf.space",
-                             port=None, api_key=QDRANT_API_KEY, timeout=10)
+qdrant_client = QdrantClient(
+    url="https://davanstrien-qdrant-test.hf.space",
+    port=None,
+    api_key=QDRANT_API_KEY,
+    timeout=10,
+)
 
 collection_name = "song_sheets"  # Replace with your actual collection name
 
+
 def search_images_by_text(query_text, top_k=5):
     # Process and encode the text query
     with torch.no_grad():
-        batch_query = colpali_processor.process_queries([query_text]).to(colpali_model.device)
+        batch_query = colpali_processor.process_queries([query_text]).to(
+            colpali_model.device
+        )
         query_embedding = colpali_model(**batch_query)
 
     # Convert the query embedding to a list of vectors
@@ -45,42 +67,101 @@ def search_images_by_text(query_text, top_k=5):
 
     return search_result
 
+
 def modify_iiif_url(url, size_percent):
     # Modify the IIIF URL to use percentage scaling
-    parts = url.split('/')
+    parts = url.split("/")
     size_index = -3
     parts[size_index] = f"pct:{size_percent}"
-    return '/'.join(parts)
+    return "/".join(parts)
+
+
+async def fetch_image(session, url):
+    async with session.get(url) as response:
+        content = await response.read()
+        return Image.open(BytesIO(content)).convert("RGB")
+
 
-def search_and_display(query, top_k, size_percent):
+async def fetch_all_images(urls):
+    async with aiohttp.ClientSession() as session:
+        tasks = [fetch_image(session, url) for url in urls]
+        return await asyncio.gather(*tasks)
+
+
+async def search_and_display(query, top_k, size_percent):
     results = search_images_by_text(query, top_k)
-    images = []
-    captions = []
-
-    for result in results.points:
-        modified_url = modify_iiif_url(result.payload['image_url'], size_percent)
-        response = requests.get(modified_url)
-        img = Image.open(BytesIO(response.content)).convert("RGB")
-        images.append(img)
-        captions.append(f"Score: {result.score:.2f}")
-
-    return images, captions
-
-# Define Gradio interface
-iface = gr.Interface(
-    fn=search_and_display,
-    inputs=[
-        gr.Textbox(label="Search Query"),
-        gr.Slider(minimum=1, maximum=20, step=1, label="Number of Results", value=5),
-        gr.Slider(minimum=1, maximum=100, step=1, label="Image Size (%)", value=100)
-    ],
-    outputs=[
-        gr.Gallery(label="Search Results", show_label=False, columns=5, height="auto"),
-        gr.JSON(label="Captions")
-    ],
-    title="Image Search with IIIF Percentage Resizing",
-    description="Enter a text query to search for images. You can adjust the number of results and the size of the returned images as a percentage of the original size."
-)
+    modified_urls = [
+        modify_iiif_url(result.payload["image_url"], size_percent)
+        for result in results.points
+    ]
+
+    images = await fetch_all_images(modified_urls)
+    html_output = (
+        "<div style='display: flex; flex-wrap: wrap; justify-content: space-around;'>"
+    )
+    for i, (image, result) in enumerate(zip(images, results.points)):
+        image_url = modified_urls[i]
+        item_url = result.payload["item_url"]
+        score = result.score
+        html_output += f"""
+        <div style='margin: 10px; text-align: center; width: 300px;'>
+            <img src='{image_url}' style='max-width: 100%; height: auto;'>
+            <p>Score: {score:.2f}</p>
+            <a href='{item_url}' target='_blank'>View Item</a>
+        </div>
+        """
+    html_output += "</div>"
+    return html_output
+
+
+# Wrapper function for synchronous Gradio interface
+def search_and_display_wrapper(query, top_k, size_percent):
+    return asyncio.run(search_and_display(query, top_k, size_percent))
+
+
+with gr.Blocks() as demo:
+    gr.HTML(
+        """
+        <h1 style='text-align: center; color: #2a4b7c;'>America Singing: Nineteenth-Century Song Sheets ColPali Search</h1>
+        <div style="display: flex; align-items: flex-start; margin-bottom: 20px;">
+            <div style="flex: 2; padding-right: 20px;">
+                <p>This app allows you to search through the Library of Congress's <a href="https://www.loc.gov/collections/nineteenth-century-song-sheets/about-this-collection/" target="_blank">"America Singing: Nineteenth-Century Song Sheets"</a> collection using natural language queries. The collection contains 4,291 song sheets from the 19th century, offering a unique window into American history, culture, and music.</p>
+
+                <p>This search functionality is powered by <a href="https://huggingface.co/blog/manu/colpali" target="_blank">ColPali</a>, an efficient document retrieval system that uses Vision Language Models. ColPali allows for searching through documents (including images and complex layouts) without the need for traditional text extraction or OCR. It works by directly embedding page images and using a <a href="https://jina.ai/news/what-is-colbert-and-late-interaction-and-why-they-matter-in-search/" target="_blank">late interaction mechanism</a> to match queries with relevant document patches.</p>
+
+                <p>ColPali's approach:
+                <ul>
+                    <li>Uses a Vision Language Model to encode document page images directly</li>
+                    <li>Splits images into patches and creates contextualized patch embeddings</li>
+                    <li>Employs a late interaction mechanism to efficiently match query tokens to document patches</li>
+                    <li>Eliminates the need for complex OCR and document parsing pipelines</li>
+                    <li>Captures both textual and visual information from documents</li>
+                </ul>
+                </p>
+            </div>
+            <div style="flex: 1;">
+                <img src="https://tile.loc.gov/image-services/iiif/service:rbc:amss:hc:00:00:3b:hc00003b:001a/full/pct:50/0/default.jpg" alt="Example Song Sheet" style="width: 100%; height: auto; border-radius: 8px; box-shadow: 0 4px 8px rgba(0,0,0,0.1);">
+                <p style="text-align: center;"><em>Example of a song sheet from the collection</em></p>
+            </div>
+        </div>
+        """
+    )
+    with gr.Row():
+        with gr.Column(scale=4):
+            search_box = gr.Textbox(
+                label="Search Query", placeholder="i.e. Irish migrant experience"
+            )
+        with gr.Column(scale=1):
+            submit_button = gr.Button("Search", variant="primary")
+    num_results = gr.Slider(
+        minimum=1, maximum=20, step=1, label="Number of Results", value=5
+    )
+    results_html = gr.HTML(label="Search Results")
+
+    submit_button.click(
+        fn=lambda query, top_k: search_and_display_wrapper(query, top_k, 100),
+        inputs=[search_box, num_results],
+        outputs=results_html,
+    )
 
-# Launch the Gradio interface
-iface.launch()
+demo.launch()