quick patch

app.py

@@ -17,10 +17,11 @@ tokenizer = BertTokenizer.from_pretrained(config.roberta_model)
 
 # --- Model Loading ---
 def load_models():
-    """Load all three models: MMRM, Text Baseline, Image Baseline."""
+    """Load models. Textual baseline is now loaded from HF Hub."""
     models = {}
+    from transformers import AutoModelForMaskedLM
     
-    # 1. MMRM
+    # 1. MMRM - Leave to None for now per user request
     # print("Loading MMRM...")
     # try:
     #     mmrm = MMRM(config).to(device)
@@ -38,44 +39,21 @@ def load_models():
     #     print(f"Error loading MMRM: {e}")
     #     models['mmrm'] = None
     models['mmrm'] = None
-
-    # 2. Textual Baseline (Fine-tuned RoBERTa)
-    print("Loading Textual Baseline...")
+    
+    # 2. Textual Baseline (Fine-tuned RoBERTa) - MIGRATED TO HF HUB
+    print("Loading Textual Baseline from HF Hub (rexera/mmrm-roberta)...")
     try:
-        # Phase 1 uses fine_tuned=True structure
-        lm_model = BaselineLanguageModel(config, fine_tuned=True).to(device)
-        # ckpt_path = config.get_phase1_checkpoint_path()
-        ckpt_path = 'rexera/mmrm-roberta'
-        if os.path.exists(ckpt_path):
-            checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)
-            
-            # Phase 1 saves 'model_state_dict' (encoder) and 'decoder_state_dict' (decoder) separately
-            # We need to map them to BaselineLanguageModel's structure: 'context_encoder' and 'classifier'
-            new_state_dict = {}
-            
-            # Map Context Encoder
-            if 'model_state_dict' in checkpoint:
-                for k, v in checkpoint['model_state_dict'].items():
-                    new_state_dict[f'context_encoder.{k}'] = v
-            
-            # Map Decoder (Classifier)
-            if 'decoder_state_dict' in checkpoint:
-                for k, v in checkpoint['decoder_state_dict'].items():
-                    new_state_dict[f'classifier.{k}'] = v
-                    
-            lm_model.load_state_dict(new_state_dict)
-            lm_model.eval()
-            models['text_baseline'] = lm_model
-            print(f"Text Baseline loaded from {ckpt_path}")
-        else:
-            print(f"Text Baseline checkpoint not found at {ckpt_path}")
-            models['text_baseline'] = None
+        # Since this is now in standard HF format (RobertaForMaskedLM)
+        repo_id = "rexera/mmrm-roberta"
+        lm_model = AutoModelForMaskedLM.from_pretrained(repo_id).to(device)
+        lm_model.eval()
+        models['text_baseline'] = lm_model
+        print(f"Textual Baseline loaded from {repo_id}")
     except Exception as e:
-        print(f"Error loading Text Baseline: {e}")
+        print(f"Error loading Textual Baseline from HF: {e}")
         models['text_baseline'] = None
-    models['text_baseline'] = None
 
-    # 3. Visual Baseline (ResNet)
+    # 3. Visual Baseline (ResNet) - Leave to None for now
     # print("Loading Visual Baseline...")
     # try:
     #     img_model = BaselineImageModel(config).to(device)
@@ -87,7 +65,7 @@ def load_models():
     #         models['visual_baseline'] = img_model
     #         print(f"Visual Baseline loaded from {ckpt_path}")
     #     else:
-    #         print(f"Visual Baseline checkpoint not found at {ckpt_path} (Expected if not deployed yet)")
+    #         print(f"Visual Baseline checkpoint not found at {ckpt_path}")
     #         models['visual_baseline'] = None
     # except Exception as e:
     #     print(f"Error loading Visual Baseline: {e}")
@@ -189,17 +167,26 @@ def run_inference(sample_idx):
     else:
         zs_res = [("Model not loaded", 0.0)]
     
-    # 1. Text Baseline
+    # 1. Textual Baseline (Fine-tuned HF Model)
     text_res = []
     if MODELS['text_baseline']:
         with torch.no_grad():
-            logits = MODELS['text_baseline'](input_ids, attention_mask, mask_positions)
-            # logits: [1, num_masks, vocab_size]
-            # Take first mask
-            mask_logits = logits[:, 0, :]
-            text_res = format_top_k(mask_logits)
+            # Standard HF model returns MaskedLMOutput
+            outputs = MODELS['text_baseline'](input_ids=input_ids, attention_mask=attention_mask)
+            all_logits = outputs.logits # [batch, seq_len, vocab_size]
+            
+            # Extract logits at mask positions
+            # input_ids/mask_positions: [1, num_masks]
+            batch_size, num_masks = mask_positions.shape
+            mask_logits = torch.gather(
+                all_logits, 1, mask_positions.unsqueeze(-1).expand(-1, -1, all_logits.size(-1))
+            )  # [batch, num_masks, vocab_size]
+            
+            # Take first mask for display
+            first_mask_logits = mask_logits[:, 0, :]
+            text_res = format_top_k(first_mask_logits)
     else:
-        text_res = [("Model not loaded (custom weight specific)", 0.0)]
+        text_res = [("Model not loaded (HF Migration)", 0.0)]
 
     # 2. Visual Baseline
     visual_res = []