Spaces:
Sleeping
Sleeping
| import torch | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| from 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} |