Upload folder using huggingface_hub

models/model_loader.py

@@ -9,23 +9,18 @@ def load_embed_model(model_path: str = "nvidia/llama-nemotron-embed-vl-1b-v2"):
     
     print(f"🔄 Loading embedding model on {device}...")
     
-    # ✅ FIX: Load CONFIG from hub but CODE from local patched file
     config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    config._attn_implementation = "sdpa"
+    if hasattr(config, 'llm_config'):
+        config.llm_config._attn_implementation = "sdpa"
     
-    # Import local patched model class
-    import sys
-    import os
-    sys.path.append(os.path.join(os.path.dirname(__file__), "local_nemotron"))
-    from local_nemotron.modeling_llama_nemotron_vl import LlamaNemotronVLModel
-    
-    # Initialize model using local class
-    model = LlamaNemotronVLModel.from_pretrained(
+    # ✅ FIX: Use manual device instead of device_map="auto"
+    model = AutoModel.from_pretrained(
         model_path,
         config=config,
         torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-        trust_remote_code=False, # We are using local code now
-        low_cpu_mem_usage=True,
-        # attn_implementation="eager", # Explicitly set in __init__ patch now
+        trust_remote_code=False,
+        low_cpu_mem_usage=True,  # ✅ CPU optimization
     ).to(device).eval()
     
     print(f"✅ Embedding model loaded on {device}")
@@ -40,20 +35,8 @@ def load_rerank_model(model_path: str = "nvidia/llama-nemotron-rerank-vl-1b-v2")
     print(f"🔄 Loading reranking model on {device}...")
     
     # ✅ FIX: Use manual device instead of device_map="auto"
-    # ✅ FIX: Load CONFIG from hub but CODE from local patched file
-    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
-    
-    # Import local patched model class
-    import sys
-    import os
-    sys.path.append(os.path.join(os.path.dirname(__file__), "local_nemotron_rerank"))
-    # Rerank model usually uses ForSequenceClassification variant, checking imports
-    from local_nemotron_rerank.modeling_llama_nemotron_vl import LlamaNemotronVLForSequenceClassification
-    
-    # Initialize model using local class
-    model = LlamaNemotronVLForSequenceClassification.from_pretrained(
+    model = AutoModelForSequenceClassification.from_pretrained(
         model_path,
-        config=config,
         torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         trust_remote_code=False,
         attn_implementation="eager",