apply fix to config dict passing for inference

inference.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, ModernBertConfig
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig, ModernBertConfig
 from typing import Dict, Any
 import yaml
 import os
@@ -84,36 +84,48 @@ class SentimentInference:
             # Load from Hugging Face Hub
             print(f"[INFERENCE_LOG] Attempting to load model from HUGGING_FACE_HUB: {model_hf_repo_id}") # Logging
             
-            # 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 = {
+            hub_config_params = {
                 "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
+                "num_weighted_layers": model_yaml_cfg.get('num_weighted_layers', 6)
             }
-            print(f"[INFERENCE_LOG] HUB_LOAD: Overrides for Hub config: {hub_config_overrides}") # Logging
+            print(f"[INFERENCE_LOG] HUB_LOAD: Parameters to update Hub config: {hub_config_params}") # 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.
+                # Step 1: Load config from Hub, allowing for our custom ModernBertConfig
+                config = ModernBertConfig.from_pretrained(model_hf_repo_id)
+                # Step 2: Update the loaded config with our specific parameters
+                for key, value in hub_config_params.items():
+                    setattr(config, key, value)
+                print(f"[INFERENCE_LOG] HUB_LOAD: Updated config: {config.to_diff_dict()}")
+
+                # Step 3: Load model with the updated config
                 self.model = ModernBertForSentiment.from_pretrained(
                     model_hf_repo_id,
-                    **hub_config_overrides
+                    config=config
                 )
-                print(f"[INFERENCE_LOG] HUB_LOAD: Model ModernBertForSentiment loaded successfully from {model_hf_repo_id}.") # Logging
+                print(f"[INFERENCE_LOG] HUB_LOAD: Model ModernBertForSentiment loaded successfully from {model_hf_repo_id} with updated config.") # 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: Error loading ModernBertForSentiment from {model_hf_repo_id} with explicit config: {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
+                
+                # Fallback: Try with AutoModelForSequenceClassification
+                # Load its config (could be BertConfig or ModernBertConfig if auto-detected)
+                # AutoConfig should ideally resolve to ModernBertConfig if architectures field is set in Hub's config.json
+                try:
+                    config_fallback = AutoConfig.from_pretrained(model_hf_repo_id)
+                    for key, value in hub_config_params.items():
+                        setattr(config_fallback, key, value)
+                    print(f"[INFERENCE_LOG] HUB_LOAD_FALLBACK: Updated fallback config: {config_fallback.to_diff_dict()}")
+
+                    self.model = AutoModelForSequenceClassification.from_pretrained(
+                        model_hf_repo_id,
+                        config=config_fallback
+                    )
+                    print(f"[INFERENCE_LOG] HUB_LOAD_FALLBACK: AutoModelForSequenceClassification loaded for {model_hf_repo_id} with updated config.") # Logging
+                except Exception as e_fallback:
+                    print(f"[INFERENCE_LOG] HUB_LOAD_FALLBACK: Critical error during fallback load: {e_fallback}")
+                    raise e_fallback # Re-raise if fallback also fails catastrophically
 
         self.model.eval()