LM ft

app.py

@@ -40,38 +40,39 @@ def load_models():
     models['mmrm'] = None
 
     # 2. Textual Baseline (Fine-tuned RoBERTa)
-    # print("Loading Textual Baseline...")
-    # try:
-    #     # Phase 1 uses fine_tuned=True structure
-    #     lm_model = BaselineLanguageModel(config, fine_tuned=True).to(device)
-    #     ckpt_path = config.get_phase1_checkpoint_path()
-    #     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
-    # except Exception as e:
-    #     print(f"Error loading Text Baseline: {e}")
-    #     models['text_baseline'] = None
+    print("Loading Textual Baseline...")
+    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
+    except Exception as e:
+        print(f"Error loading Text Baseline: {e}")
+        models['text_baseline'] = None
     models['text_baseline'] = None
 
     # 3. Visual Baseline (ResNet)