Refactor app.py: improve code formatting and enhance readability

app.py

@@ -23,9 +23,16 @@ import seaborn as sns
 from einops import rearrange
 
 # Import from colpali_engine
-from colpali_engine.models import BiGemma3, BiGemmaProcessor3, ColGemma3, ColGemmaProcessor3
+from colpali_engine.models import (
+    BiGemma3,
+    BiGemmaProcessor3,
+    ColGemma3,
+    ColGemmaProcessor3,
+)
 from colpali_engine.interpretability import get_similarity_maps_from_embeddings
-from colpali_engine.interpretability.similarity_map_utils import normalize_similarity_map
+from colpali_engine.interpretability.similarity_map_utils import (
+    normalize_similarity_map,
+)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -33,6 +40,7 @@ print(f"Device: {device}")
 if torch.cuda.is_available():
     print(f"GPU: {torch.cuda.get_device_name(0)}")
 
+
 # Global state for models and indexed documents
 class DocumentIndex:
     def __init__(self):
@@ -44,6 +52,7 @@ class DocumentIndex:
         self.colgemma_model = None
         self.colgemma_processor = None
 
+
 doc_index = DocumentIndex()
 
 
@@ -113,7 +122,7 @@ def index_bigemma_images(images: List[Image.Image]):
     # Process in smaller batches to avoid memory issues
     batch_size = 2
     for i in range(0, len(images), batch_size):
-        batch = images[i:i+batch_size]
+        batch = images[i : i + batch_size]
         batch_images = processor.process_images(batch).to(device)
 
         with torch.no_grad():
@@ -122,7 +131,9 @@ def index_bigemma_images(images: List[Image.Image]):
 
     # Concatenate all embeddings
     all_embeddings = torch.cat(embeddings_list, dim=0)
-    print(f"✓ Indexed {len(images)} pages with BiGemma3 (shape: {all_embeddings.shape})")
+    print(
+        f"✓ Indexed {len(images)} pages with BiGemma3 (shape: {all_embeddings.shape})"
+    )
 
     return all_embeddings
 
@@ -138,7 +149,7 @@ def index_colgemma_images(images: List[Image.Image]):
     # Process in smaller batches to avoid memory issues
     batch_size = 2
     for i in range(0, len(images), batch_size):
-        batch = images[i:i+batch_size]
+        batch = images[i : i + batch_size]
         batch_images = processor.process_images(batch).to(device)
 
         with torch.no_grad():
@@ -147,7 +158,9 @@ def index_colgemma_images(images: List[Image.Image]):
 
     # Concatenate all embeddings
     all_embeddings = torch.cat(embeddings_list, dim=0)
-    print(f"✓ Indexed {len(images)} pages with ColGemma3 (shape: {all_embeddings.shape})")
+    print(
+        f"✓ Indexed {len(images)} pages with ColGemma3 (shape: {all_embeddings.shape})"
+    )
 
     return all_embeddings
 
@@ -182,11 +195,14 @@ def index_document(pdf_files, model_choice: str) -> str:
             doc_index.colgemma_embeddings = index_colgemma_images(doc_index.images)
             status_messages.append("✓ Indexed with ColGemma3")
 
-        final_status = "\n".join(status_messages) + "\n\n✅ Document ready for querying!"
+        final_status = (
+            "\n".join(status_messages) + "\n\n✅ Document ready for querying!"
+        )
         return final_status
 
     except Exception as e:
         import traceback
+
         error_details = traceback.format_exc()
         print(f"Indexing error: {error_details}")
         return f"❌ Error indexing document: {str(e)}"
@@ -211,14 +227,18 @@ def generate_colgemma_heatmap(
             image_mask = batch_images["input_ids"] == image_token_id
         else:
             image_mask = torch.ones(
-                image_embedding.shape[0], image_embedding.shape[1],
-                dtype=torch.bool, device=device
+                image_embedding.shape[0],
+                image_embedding.shape[1],
+                dtype=torch.bool,
+                device=device,
             )
 
         # Calculate n_patches
         num_image_tokens = image_mask.sum().item()
         n_side = int(math.sqrt(num_image_tokens))
-        n_patches = (n_side, n_side) if n_side * n_side == num_image_tokens else (16, 16)
+        n_patches = (
+            (n_side, n_side) if n_side * n_side == num_image_tokens else (16, 16)
+        )
 
         # Generate similarity maps
         similarity_maps_list = get_similarity_maps_from_embeddings(
@@ -235,12 +255,14 @@ def generate_colgemma_heatmap(
 
         # Create heatmap overlay
         img_array = np.array(image.convert("RGBA"))
-        similarity_map_array = normalize_similarity_map(aggregated_map).to(torch.float32).cpu().numpy()
+        similarity_map_array = (
+            normalize_similarity_map(aggregated_map).to(torch.float32).cpu().numpy()
+        )
         similarity_map_array = rearrange(similarity_map_array, "h w -> w h")
 
-        similarity_map_image = Image.fromarray((similarity_map_array * 255).astype("uint8")).resize(
-            image.size, Image.Resampling.BICUBIC
-        )
+        similarity_map_image = Image.fromarray(
+            (similarity_map_array * 255).astype("uint8")
+        ).resize(image.size, Image.Resampling.BICUBIC)
 
         # Create matplotlib figure
         fig, ax = plt.subplots(figsize=(10, 10))
@@ -287,7 +309,12 @@ def query_documents(
         # Query with BiGemma3
         if model_choice in ["NetraEmbed (BiGemma3)", "Both"]:
             if doc_index.bigemma_embeddings is None:
-                return None, "⚠️ Please index the document with BiGemma3 first.", None, None
+                return (
+                    None,
+                    "⚠️ Please index the document with BiGemma3 first.",
+                    None,
+                    None,
+                )
 
             model, processor = load_bigemma_model()
 
@@ -311,15 +338,27 @@ def query_documents(
             bigemma_text = "### BiGemma3 (NetraEmbed) Results\n\n"
             for rank, idx in enumerate(top_indices):
                 score = scores[0, idx].item()
-                bigemma_text += f"**Rank {rank + 1}:** Page {idx.item() + 1} - Score: {score:.4f}\n"
+                bigemma_text += (
+                    f"**Rank {rank + 1}:** Page {idx.item() + 1} - Score: {score:.4f}\n"
+                )
                 bigemma_results.append(
-                    (doc_index.images[idx.item()], f"Rank {rank + 1} - Page {idx.item() + 1} (Score: {score:.4f})")
+                    (
+                        doc_index.images[idx.item()],
+                        f"Rank {rank + 1} - Page {idx.item() + 1} (Score: {score:.4f})",
+                    )
                 )
 
         # Query with ColGemma3
         if model_choice in ["ColNetraEmbed (ColGemma3)", "Both"]:
             if doc_index.colgemma_embeddings is None:
-                return bigemma_results if bigemma_results else None, bigemma_text if bigemma_text else "⚠️ Please index the document with ColGemma3 first.", None, None
+                return (
+                    bigemma_results if bigemma_results else None,
+                    bigemma_text
+                    if bigemma_text
+                    else "⚠️ Please index the document with ColGemma3 first.",
+                    None,
+                    None,
+                )
 
             model, processor = load_colgemma_model()
 
@@ -343,7 +382,9 @@ def query_documents(
             colgemma_text = "### ColGemma3 (ColNetraEmbed) Results\n\n"
             for rank, idx in enumerate(top_indices):
                 score = scores[0, idx].item()
-                colgemma_text += f"**Rank {rank + 1}:** Page {idx.item() + 1} - Score: {score:.2f}\n"
+                colgemma_text += (
+                    f"**Rank {rank + 1}:** Page {idx.item() + 1} - Score: {score:.2f}\n"
+                )
 
                 # Generate heatmap if requested
                 if show_heatmap:
@@ -353,11 +394,17 @@ def query_documents(
                         image_embedding=doc_index.colgemma_embeddings[idx.item()],
                     )
                     colgemma_results.append(
-                        (heatmap_image, f"Rank {rank + 1} - Page {idx.item() + 1} (Score: {score:.2f})")
+                        (
+                            heatmap_image,
+                            f"Rank {rank + 1} - Page {idx.item() + 1} (Score: {score:.2f})",
+                        )
                     )
                 else:
                     colgemma_results.append(
-                        (doc_index.images[idx.item()], f"Rank {rank + 1} - Page {idx.item() + 1} (Score: {score:.2f})")
+                        (
+                            doc_index.images[idx.item()],
+                            f"Rank {rank + 1} - Page {idx.item() + 1} (Score: {score:.2f})",
+                        )
                     )
 
         # Return results based on model choice
@@ -370,6 +417,7 @@ def query_documents(
 
     except Exception as e:
         import traceback
+
         error_details = traceback.format_exc()
         print(f"Query error: {error_details}")
         return None, f"❌ Error during query: {str(e)}", None, None
@@ -390,14 +438,14 @@ with gr.Blocks(title="NetraEmbed Demo") as demo:
                     <a href="https://github.com/adithya-s-k/colpali" target="_blank">
                         <img src="https://img.shields.io/badge/GitHub-colpali-181717?logo=github" alt="GitHub">
                     </a>
-                    <a href="https://huggingface.co/Cognitive-Lab/ColNetraEmbed" target="_blank">
+                    <a href="https://huggingface.co/Cognitive-Lab/NetraEmbed" target="_blank">
                         <img src="https://img.shields.io/badge/🤗%20HuggingFace-Model-yellow" alt="Model">
                     </a>
                     <a href="https://www.cognitivelab.in/blog/introducing-netraembed" target="_blank">
                         <img src="https://img.shields.io/badge/Blog-CognitiveLab-blue" alt="Blog">
                     </a>
-                    <a href="https://cloud.cognitivelab.in" target="_blank">
-                        <img src="https://img.shields.io/badge/Demo-Try%20it%20out-green" alt="Demo">
+                    <a href="https://huggingface.co/spaces/AdithyaSK/NetraEmbed" target="_blank">
+                        <img src="https://img.shields.io/badge/🤗%20Demo-HuggingFace%20Space-yellow" alt="Demo">
                     </a>
                 </div>
                 """
@@ -443,9 +491,7 @@ with gr.Blocks(title="NetraEmbed Demo") as demo:
             )
 
             pdf_upload = gr.File(
-                label="Upload PDFs",
-                file_types=[".pdf"],
-                file_count="multiple"
+                label="Upload PDFs", file_types=[".pdf"], file_count="multiple"
             )
             index_btn = gr.Button("📥 Index Documents", variant="primary", size="sm")
 
@@ -531,7 +577,12 @@ with gr.Blocks(title="NetraEmbed Demo") as demo:
     query_btn.click(
         fn=query_documents,
         inputs=[query_input, model_select, top_k_slider, heatmap_checkbox],
-        outputs=[bigemma_gallery, bigemma_results_text, colgemma_gallery, colgemma_results_text],
+        outputs=[
+            bigemma_gallery,
+            bigemma_results_text,
+            colgemma_gallery,
+            colgemma_results_text,
+        ],
     )
 
 # Enable queue for handling multiple requests