Update rag/models.py

rag/models.py

@@ -6,57 +6,54 @@ from transformers import (
     Qwen2VLForConditionalGeneration,
     AutoProcessor,
 )
-
 from gliner import GLiNER
+
 from rag.config import Settings
 from rag.logging_utils import get_logger
 
 logger = get_logger(__name__)
 
-def _has_cuda() -> bool:
-    return torch.cuda.is_available()
-
-def _dtype_for_device():
-    # BF16 is great on recent GPUs; on CPU it's often slower / unsupported in ops
-    return torch.bfloat16 if _has_cuda() else torch.float32
-
 class Models:
     def __init__(self, settings: Settings):
         self.settings = settings
 
-        # Router (CPU)
-        logger.info("Loading Router (GLiNER) on CPU: %s", settings.router_model_id)
-        self.router = GLiNER.from_pretrained(settings.router_model_id).to("cpu")
-        self.router.eval()
+        # Router CPU
+        logger.info("🧠 Loading GLiNER router on CPU: %s", settings.router_model_id)
+        self.router_model = GLiNER.from_pretrained(settings.router_model_id).to("cpu")
+        self.router_model.eval()
 
-        # Embedder
-        logger.info("Loading Embedder: %s", settings.embed_model_id)
+        # Embedding
+        logger.info("🔹 Loading embedder: %s", settings.embed_model_id)
         self.embed_tokenizer = AutoTokenizer.from_pretrained(settings.embed_model_id, trust_remote_code=False)
         self.embed_model = AutoModel.from_pretrained(
             settings.embed_model_id,
             trust_remote_code=False,
-            torch_dtype=_dtype_for_device(),
+            torch_dtype=torch.bfloat16,
             device_map="auto",
         )
         self.embed_model.eval()
 
         # Reranker
-        logger.info("Loading Reranker: %s", settings.rerank_model_id)
+        logger.info("⚖️ Loading reranker: %s", settings.rerank_model_id)
         self.rerank_tokenizer = AutoTokenizer.from_pretrained(settings.rerank_model_id)
         self.rerank_model = AutoModelForSequenceClassification.from_pretrained(
             settings.rerank_model_id,
-            torch_dtype=_dtype_for_device(),
+            torch_dtype=torch.bfloat16,
             device_map="auto",
         )
         self.rerank_model.eval()
 
-        # Vision generator
-        logger.info("Loading Vision model: %s", settings.gen_model_id)
+        # Vision
+        logger.info("👁️ Loading vision model: %s", settings.gen_model_id)
         self.gen_model = Qwen2VLForConditionalGeneration.from_pretrained(
             settings.gen_model_id,
-            torch_dtype=_dtype_for_device(),
+            torch_dtype=torch.bfloat16,
             device_map="auto",
         )
         self.gen_model.eval()
-
         self.gen_processor = AutoProcessor.from_pretrained(settings.gen_model_id)
+
+def load_models(settings: Settings | None = None) -> Models:
+    if settings is None:
+        settings = Settings()
+    return Models(settings)