Upload src/inference.py with huggingface_hub

src/inference.py

@@ -0,0 +1,79 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from src.models import ModernBertForSentiment
+from transformers import ModernBertConfig
+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."""
+        with open(config_path, 'r') as f:
+            config = yaml.safe_load(f)
+        
+        model_cfg = config.get('model', {})
+        inference_cfg = config.get('inference', {})
+        
+        # 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'))
+        
+        # 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))
+
+        # 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
+
+        # 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
+
+        print("Instantiating ModernBertForSentiment model structure...")
+        self.model = ModernBertForSentiment(bert_config)
+        
+        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'))
+        
+        # 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)
+        self.model.eval()
+        print("Model loaded successfully.")
+        
+    def predict(self, text: str) -> Dict[str, Any]:
+        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=self.max_length)
+        with torch.no_grad():
+            outputs = self.model(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
+        logits = outputs["logits"]
+        prob = torch.sigmoid(logits).item()
+        return {"sentiment": "positive" if prob > 0.5 else "negative", "confidence": prob}