DEFAULT_MODEL = ColSmol-500M

src/colpali/visual_search.py

@@ -9,8 +9,9 @@ All dependencies are now within src/colpali/ - no external colpali_colab_package
 
 import logging
 from typing import List, Dict, Any, Optional
-import numpy as np
+
 import torch
+import numpy as np
 from qdrant_client import QdrantClient
 
 # Import from local src/colpali modules (no external dependencies)
@@ -21,6 +22,7 @@ from src.colpali.search import VisualDocumentSearch
 from src.utils import get_device_for_colpali
 
 logger = logging.getLogger(__name__)
+DEFAULT_MODEL = "colSmol-500M"
 
 
 class VisualSearchResult:
@@ -53,8 +55,8 @@ class VisualSearchAdapter:
         self,
         qdrant_url: str,
         qdrant_api_key: str,
-        collection_name: str = "colSmol-500M",
-        model_name: str = "vidore/colSmol-500M",
+        collection_name: str = DEFAULT_MODEL,
+        model_name: str = f"vidore/{DEFAULT_MODEL}",
         device: str = None,
         batch_size: int = 4
     ):
@@ -207,7 +209,7 @@ class VisualSearchAdapter:
 def create_visual_search_adapter(
     qdrant_url: Optional[str] = None,
     qdrant_api_key: Optional[str] = None,
-    collection_name: str = "colSmol-500M"
+    collection_name: str = DEFAULT_MODEL
 ) -> VisualSearchAdapter:
     """
     Factory function to create a visual search adapter.