update inference to load from hf model id

config.yaml

@@ -1,6 +1,6 @@
 model:
   name: "voxmenthe/modernbert-imdb-sentiment"
-  output_dir: "checkpoints"
+  tokenizer_name_or_path: "answerdotai/ModernBERT-base"
   max_length: 880 # 256
   dropout: 0.1
   pooling_strategy: "mean" # Current default, change as needed

inference.py

@@ -1,79 +1,65 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from models import ModernBertForSentiment
-from transformers import ModernBertConfig
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, ModernBertConfig
+# models.py (containing ModernBertForSentiment) will be loaded from the Hub due to trust_remote_code=True
 from typing import Dict, Any
 import yaml
-import os
-
 
 class SentimentInference:
     def __init__(self, config_path: str = "config.yaml"):
-        """Load configuration and initialize model and tokenizer."""
+        """Load configuration and initialize model and tokenizer from Hugging Face Hub."""
         with open(config_path, 'r') as f:
-            config = yaml.safe_load(f)
-        
-        model_cfg = config.get('model', {})
-        inference_cfg = config.get('inference', {})
+            config_data = yaml.safe_load(f)
         
-        # Path to the .pt model weights file
-        model_weights_path = inference_cfg.get('model_path', 
-                                             os.path.join(model_cfg.get('output_dir', 'checkpoints'), 'best_model.pt'))
+        model_yaml_cfg = config_data.get('model', {})
+        inference_yaml_cfg = config_data.get('inference', {})
         
-        # Base model name from config (e.g., 'answerdotai/ModernBERT-base')
-        # This will be used for loading both tokenizer and base BERT config from Hugging Face Hub
-        base_model_name = model_cfg.get('name', 'answerdotai/ModernBERT-base')
-
-        self.max_length = inference_cfg.get('max_length', model_cfg.get('max_length', 256))
+        model_hf_repo_id = model_yaml_cfg.get('name_or_path')
+        if not model_hf_repo_id:
+            raise ValueError("model.name_or_path must be specified in config.yaml (e.g., 'username/model_name')")
 
-        # Load tokenizer from the base model name (e.g., from Hugging Face Hub)
-        print(f"Loading tokenizer from: {base_model_name}")
-        self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-        
-        # Load base BERT config from the base model name
-        print(f"Loading ModernBertConfig from: {base_model_name}")
-        bert_config = ModernBertConfig.from_pretrained(base_model_name) 
-        
-        # --- Apply any necessary overrides from your config to the loaded bert_config --- 
-        # For example, if your ModernBertForSentiment expects specific config values beyond the base BERT model.
-        # Your current ModernBertForSentiment takes the entire config object, which might implicitly carry these.
-        # However, explicitly setting them on bert_config loaded from HF is safer if they are architecturally relevant.
-        bert_config.classifier_dropout = model_cfg.get('dropout', bert_config.classifier_dropout) # Example
-        # Ensure num_labels is set if your inference model needs it (usually for HF pipeline, less so for manual predict)
-        # bert_config.num_labels = model_cfg.get('num_labels', 1) # Typically 1 for binary sentiment regression-style output
+        tokenizer_hf_repo_id = model_yaml_cfg.get('tokenizer_name_or_path', model_hf_repo_id)
 
-        # It's also important that pooling_strategy and num_weighted_layers are set on the config object 
-        # that ModernBertForSentiment receives, as it uses these to build its layers.
-        # These are usually fine-tuning specific, not part of the base HF config, so they should come from your model_cfg.
-        bert_config.pooling_strategy = model_cfg.get('pooling_strategy', 'cls')
-        bert_config.num_weighted_layers = model_cfg.get('num_weighted_layers', 4)
-        bert_config.loss_function = model_cfg.get('loss_function', {'name': 'SentimentWeightedLoss', 'params': {}}) # Needed by model init
-        # Ensure num_labels is explicitly set for the model's classifier head
-        bert_config.num_labels = 1 # For sentiment (positive/negative) often treated as 1 logit output
+        self.max_length = inference_yaml_cfg.get('max_length', model_yaml_cfg.get('max_length', 512))
 
-        print("Instantiating ModernBertForSentiment model structure...")
-        self.model = ModernBertForSentiment(bert_config)
+        print(f"Loading tokenizer from: {tokenizer_hf_repo_id}")
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_hf_repo_id)
         
-        print(f"Loading model weights from local checkpoint: {model_weights_path}")
-        # Load the entire checkpoint dictionary first
-        checkpoint = torch.load(model_weights_path, map_location=torch.device('cpu'))
+        print(f"Loading base ModernBertConfig from: {model_hf_repo_id}")
+        # Load the config that was uploaded with the model (config.json in the HF repo)
+        # This config should already have the correct architecture defined by ModernBertConfig.
+        # We then augment it with any custom parameters needed by ModernBertForSentiment's __init__.
+        loaded_config = ModernBertConfig.from_pretrained(model_hf_repo_id)
         
-        # Extract the model_state_dict from the checkpoint
-        # This handles the case where the checkpoint saves more than just the model weights (e.g., optimizer state, epoch)
-        if 'model_state_dict' in checkpoint:
-            model_state_to_load = checkpoint['model_state_dict']
-        else:
-            # If the checkpoint is just the state_dict itself (older format or different saving convention)
-            model_state_to_load = checkpoint
-            
-        self.model.load_state_dict(model_state_to_load)
+        # Augment loaded_config with parameters from model_yaml_cfg needed for ModernBertForSentiment initialization
+        # These should reflect how the model was trained and its specific custom head.
+        loaded_config.pooling_strategy = model_yaml_cfg.get('pooling_strategy', 'mean') # Default to 'mean' as per your models.py change
+        loaded_config.num_weighted_layers = model_yaml_cfg.get('num_weighted_layers', 4)
+        loaded_config.classifier_dropout = model_yaml_cfg.get('dropout') # Allow None if not in yaml
+        # num_labels should ideally be in the config.json uploaded to HF, but can be set here if needed.
+        # For binary sentiment with a single logit output, num_labels is 1.
+        loaded_config.num_labels = model_yaml_cfg.get('num_labels', 1)
+        # The loss_function might not be strictly needed for inference if the model doesn't use it in forward pass for eval,
+        # but if ModernBertForSentiment.__init__ requires it, it must be provided.
+        # Assuming it's not critical for basic inference here to simplify.
+        # loaded_config.loss_function = model_yaml_cfg.get('loss_function', {'name': '...', 'params': {}})
+
+        print(f"Instantiating and loading model weights for {model_hf_repo_id}...")
+        # trust_remote_code=True allows loading models.py (containing ModernBertForSentiment)
+        # from the Hugging Face model repository.
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            model_hf_repo_id,
+            config=loaded_config, # Pass the augmented config
+            trust_remote_code=True
+        )
         self.model.eval()
-        print("Model loaded successfully.")
+        print(f"Model {model_hf_repo_id} loaded successfully from Hugging Face Hub.")
         
     def predict(self, text: str) -> Dict[str, Any]:
-        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=self.max_length)
+        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=self.max_length, padding=True)
         with torch.no_grad():
             outputs = self.model(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
-        logits = outputs["logits"]
+        logits = outputs.get("logits") # Use .get for safety
+        if logits is None:
+            raise ValueError("Model output did not contain 'logits'. Check model's forward pass.")
         prob = torch.sigmoid(logits).item()
         return {"sentiment": "positive" if prob > 0.5 else "negative", "confidence": prob}