Update pipeline/__init__.py

pipeline/__init__.py

@@ -1,6 +1,7 @@
 import os
 import torch
 from flask import Flask
+from sentence_transformers.cross_encoder import CrossEncoder 
 
 FEATURE_WEIGHTS = {"shape": 0.4, "color": 0.5, "texture": 0.1}
 FINAL_SCORE_THRESHOLD = 0.5
@@ -24,7 +25,6 @@ processor_gnd = AutoProcessor.from_pretrained(gnd_model_id)
 model_gnd = AutoModelForZeroShotObjectDetection.from_pretrained(gnd_model_id).to(device)
 
 print("...Loading Segment Anything (SAM) model...")
-# IMPORTANT: The path is now relative to the root of the project
 sam_checkpoint = "sam_vit_b_01ec64.pth"
 sam_model = sam_model_registry["vit_b"](checkpoint=sam_checkpoint).to(device)
 predictor = SamPredictor(sam_model)
@@ -34,7 +34,7 @@ bge_model_id = "BAAI/bge-small-en-v1.5"
 tokenizer_text = AutoTokenizer.from_pretrained(bge_model_id)
 model_text = AutoModel.from_pretrained(bge_model_id).to(device)
 
-# Store models in a dictionary to pass to logic functions
+
 models = {
     "processor_gnd": processor_gnd,
     "model_gnd": model_gnd,
@@ -44,8 +44,16 @@ models = {
     "device": device
 }
 
+
+print("...Loading Cross-Encoder model for re-ranking...")
+
+cross_encoder_model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device=device)
+
+
+models["cross_encoder"] = cross_encoder_model
+
+
 print("✅ All models loaded successfully.")
 print("="*50)
 
-# Import routes after app and models are defined to avoid circular imports
 from pipeline import routes