Update app.py

app.py

@@ -11,7 +11,6 @@ from pdf2image import convert_from_path
 import torch
 from torch.utils.data import DataLoader
 from transformers.utils.import_utils import is_flash_attn_2_available
-# from colpali_engine.models import ColQwen2, ColQwen2Processor
 from colpali_engine.models import ColQwen2_5, ColQwen2_5_Processor
 from openai import OpenAI
 import spaces
@@ -21,12 +20,11 @@ import gradio as gr
 # Enable flash attention
 # subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-# Load the visual document retrieval model (ColQwen2)
+# Load the visual document retrieval model
 model = ColQwen2_5.from_pretrained(
     "vidore/colqwen2.5-v0.2",
     torch_dtype=torch.bfloat16,
     device_map="cuda:0",
-    # attn_implementation="flash_attention_2",
     attn_implementation="flash_attention_2" if is_flash_attn_2_available() else None,
 ).eval()
 processor = ColQwen2_5_Processor.from_pretrained("vidore/colqwen2.5-v0.2")
@@ -59,7 +57,7 @@ def convert_files(files):
 ################################################
 @spaces.GPU
 def index_gpu(images, ds):
-    """Runs inference on the GPU for the given images with ColQwen model."""
+    """Runs inference on the GPU for the given images with the visual document retrieval model."""
     # Specify the device
     device = "cuda:0" if torch.cuda.is_available() else "cpu"
     if device != model.device: