new test inference

inference.py

@@ -29,78 +29,92 @@ class SentimentInference:
         if not tokenizer_hf_repo_id and not model_hf_repo_id:
             raise ValueError("Either model.tokenizer_name_or_path or model.name_or_path (as fallback for tokenizer) must be specified in config.yaml")
         effective_tokenizer_repo_id = tokenizer_hf_repo_id or model_hf_repo_id
-        print(f"Loading tokenizer from: {effective_tokenizer_repo_id}")
+        print(f"[INFERENCE_LOG] Loading tokenizer from: {effective_tokenizer_repo_id}") # Logging
         self.tokenizer = AutoTokenizer.from_pretrained(effective_tokenizer_repo_id)
 
         # --- Model Loading --- #
         # Determine if we are loading from a local .pt file or from Hugging Face Hub
         load_from_local_pt = False
         if local_model_weights_path and os.path.isfile(local_model_weights_path):
-            print(f"Found local model weights path: {local_model_weights_path}")
+            print(f"[INFERENCE_LOG] Found local model weights path: {local_model_weights_path}") # Logging
             print(f"--- Debug: Found local model weights path: {local_model_weights_path} ---") # Add this
             load_from_local_pt = True
         elif not model_hf_repo_id:
             raise ValueError("No local model_path found and model.name_or_path (for Hub) is not specified in config.yaml")
 
+        print(f"[INFERENCE_LOG] load_from_local_pt: {load_from_local_pt}") # Logging
         print(f"--- Debug: load_from_local_pt is: {load_from_local_pt} ---") # Add this
 
         if load_from_local_pt:
-            print("Attempting to load model from local .pt checkpoint...")
+            print("[INFERENCE_LOG] Attempting to load model from LOCAL .pt checkpoint...") # Logging
             print("--- Debug: Entering LOCAL .pt loading path ---") # Add this
             # Base BERT config must still be loaded, usually from a Hub ID (e.g., original base model)
             # This base_model_for_config_id is crucial for building the correct ModernBertForSentiment structure.
-            base_model_for_config_id = model_yaml_cfg.get('base_model_for_config', model_hf_repo_id or tokenizer_hf_repo_id)
-            print(f"--- Debug: base_model_for_config_id (for local .pt): {base_model_for_config_id} ---") # Add this
+            base_model_for_config_id = model_yaml_cfg.get('base_model_for_config', model_yaml_cfg.get('name_or_path'))
             if not base_model_for_config_id:
-                 raise ValueError("For local .pt loading, model.base_model_for_config must be specified in config.yaml (e.g., 'answerdotai/ModernBERT-base') to build the model structure.")
+                 raise ValueError("model.base_model_for_config or model.name_or_path must be specified in config.yaml when loading local .pt for ModernBertForSentiment structure.")
             
-            print(f"Loading ModernBertConfig for structure from: {base_model_for_config_id}")
-            bert_config = ModernBertConfig.from_pretrained(base_model_for_config_id)
-            
-            # Augment config with parameters from model_yaml_cfg
-            bert_config.pooling_strategy = model_yaml_cfg.get('pooling_strategy', 'mean')
-            bert_config.num_weighted_layers = model_yaml_cfg.get('num_weighted_layers', 4)
-            bert_config.classifier_dropout = model_yaml_cfg.get('dropout')
-            bert_config.num_labels = model_yaml_cfg.get('num_labels', 1)
-            # bert_config.loss_function = model_yaml_cfg.get('loss_function') # If needed by __init__
+            print(f"[INFERENCE_LOG] LOCAL_PT_LOAD: base_model_for_config_id: {base_model_for_config_id}") # Logging
+
+            model_config = ModernBertConfig.from_pretrained(
+                base_model_for_config_id, 
+                num_labels=model_yaml_cfg.get('num_labels', 1), # from config.yaml via model_yaml_cfg
+                pooling_strategy=model_yaml_cfg.get('pooling_strategy', 'mean'), # from config.yaml via model_yaml_cfg
+                num_weighted_layers=model_yaml_cfg.get('num_weighted_layers', 4) # from config.yaml via model_yaml_cfg
+            )
+            print(f"[INFERENCE_LOG] LOCAL_PT_LOAD: Loaded ModernBertConfig: {model_config.to_diff_dict()}") # Logging
 
-            print("Instantiating ModernBertForSentiment model structure...")
-            self.model = ModernBertForSentiment(bert_config)
+            print(f"[INFERENCE_LOG] LOCAL_PT_LOAD: Initializing ModernBertForSentiment with this config.") # Logging
+            self.model = ModernBertForSentiment(config=model_config)
             
-            print(f"Loading model weights from local checkpoint: {local_model_weights_path}")
+            print(f"[INFERENCE_LOG] LOCAL_PT_LOAD: Loading weights from checkpoint: {local_model_weights_path}") # Logging
             checkpoint = torch.load(local_model_weights_path, map_location=torch.device('cpu'))
-            if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
-                model_state_to_load = checkpoint['model_state_dict']
-            else:
-                model_state_to_load = checkpoint # Assume it's the state_dict itself
-            self.model.load_state_dict(model_state_to_load)
-            print(f"Model loaded successfully from local checkpoint: {local_model_weights_path}.")
+            
+            state_dict_to_load = checkpoint.get('model_state_dict', checkpoint.get('state_dict', checkpoint))
+            if not isinstance(state_dict_to_load, dict):
+                raise TypeError(f"Loaded checkpoint from {local_model_weights_path} is not a dict or does not contain 'model_state_dict' or 'state_dict'.")
 
-        else: # Load from Hugging Face Hub
-            print(f"Attempting to load model from Hugging Face Hub: {model_hf_repo_id}...")
-            print(f"--- Debug: Entering HUGGING FACE HUB loading path ---") # Add this
-            print(f"--- Debug: model_hf_repo_id (for Hub loading): {model_hf_repo_id} ---") # Add this
-            if not model_hf_repo_id:
-                raise ValueError("model.name_or_path must be specified in config.yaml for Hub loading.")
+            # Log first few keys for debugging
+            first_few_keys = list(state_dict_to_load.keys())[:5]
+            print(f"[INFERENCE_LOG] LOCAL_PT_LOAD: First few keys from checkpoint state_dict: {first_few_keys}") # Logging
 
-            print(f"Loading base ModernBertConfig from: {model_hf_repo_id}")
-            loaded_config = ModernBertConfig.from_pretrained(model_hf_repo_id)
+            self.model.load_state_dict(state_dict_to_load)
+            print(f"[INFERENCE_LOG] LOCAL_PT_LOAD: Weights loaded successfully into ModernBertForSentiment from {local_model_weights_path}.") # Logging
+        else:
+            # Load from Hugging Face Hub
+            print(f"[INFERENCE_LOG] Attempting to load model from HUGGING_FACE_HUB: {model_hf_repo_id}") # Logging
             
-            # Augment loaded_config
-            loaded_config.pooling_strategy = model_yaml_cfg.get('pooling_strategy', 'mean')
-            loaded_config.num_weighted_layers = model_yaml_cfg.get('num_weighted_layers', 6) # Default to 6 now
-            loaded_config.classifier_dropout = model_yaml_cfg.get('dropout')
-            loaded_config.num_labels = model_yaml_cfg.get('num_labels', 1)
+            # Here, we use the config that's packaged with the model on the Hub by default.
+            # We just add/override num_labels, pooling_strategy, num_weighted_layers if they are in our local config.yaml
+            # as these might be specific to our fine-tuning and not in the Hub's default config.json.
+            hub_config_overrides = {
+                "num_labels": model_yaml_cfg.get('num_labels', 1),
+                "pooling_strategy": model_yaml_cfg.get('pooling_strategy', 'mean'),
+                "num_weighted_layers": model_yaml_cfg.get('num_weighted_layers', 6) # Default to 6 now
+            }
+            print(f"[INFERENCE_LOG] HUB_LOAD: Overrides for Hub config: {hub_config_overrides}") # Logging
+
+            try:
+                # Using ModernBertForSentiment.from_pretrained directly.
+                # This assumes the config.json on the Hub for 'model_hf_repo_id' is compatible
+                # or that from_pretrained can correctly initialize ModernBertForSentiment with it.
+                self.model = ModernBertForSentiment.from_pretrained(
+                    model_hf_repo_id,
+                    **hub_config_overrides
+                )
+                print(f"[INFERENCE_LOG] HUB_LOAD: Model ModernBertForSentiment loaded successfully from {model_hf_repo_id}.") # Logging
+            except Exception as e:
+                print(f"[INFERENCE_LOG] HUB_LOAD: Error loading ModernBertForSentiment from {model_hf_repo_id}: {e}") # Logging
+                print(f"[INFERENCE_LOG] HUB_LOAD: Falling back to AutoModelForSequenceClassification for {model_hf_repo_id}.") # Logging
+                # Fallback: Try with AutoModelForSequenceClassification if ModernBertForSentiment fails
+                # This might happen if the Hub model isn't strictly saved as a ModernBertForSentiment type
+                # or if its config.json doesn't have _custom_class set, etc.
+                self.model = AutoModelForSequenceClassification.from_pretrained(
+                    model_hf_repo_id,
+                    **hub_config_overrides
+                )
+                print(f"[INFERENCE_LOG] HUB_LOAD: AutoModelForSequenceClassification loaded for {model_hf_repo_id}.") # Logging
 
-            print(f"Instantiating and loading model weights for {model_hf_repo_id}...")
-            self.model = AutoModelForSequenceClassification.from_pretrained(
-                model_hf_repo_id,
-                config=loaded_config,
-                trust_remote_code=True,
-                force_download=True  # <--- TEMPORARY - remove when everything is working
-            )
-            print(f"Model {model_hf_repo_id} loaded successfully from Hugging Face Hub.")
-        
         self.model.eval()
         
     def predict(self, text: str) -> Dict[str, Any]: