Create models.py

rag/models.py

@@ -0,0 +1,62 @@
+import torch
+from transformers import (
+    AutoModel,
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    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()
+
+        # Embedder
+        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(),
+            device_map="auto",
+        )
+        self.embed_model.eval()
+
+        # Reranker
+        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(),
+            device_map="auto",
+        )
+        self.rerank_model.eval()
+
+        # Vision generator
+        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(),
+            device_map="auto",
+        )
+        self.gen_model.eval()
+
+        self.gen_processor = AutoProcessor.from_pretrained(settings.gen_model_id)