# This file defines the custom architecture for the hybrid model.

import torch
import torch.nn as nn
from transformers import AutoModel, PreTrainedModel, AutoConfig

class HybridRegressionModel(PreTrainedModel):
    """
    A hybrid model that combines a pre-trained transformer with a regression head
    that accepts additional numerical features.
    This architecture must exactly match the one used during training.
    """
    # This associates the model with the base configuration class
    config_class = AutoConfig

    def __init__(self, config, num_extra_features=7):
        super(HybridRegressionModel, self).__init__(config)
        # Load the transformer body from the configuration
        self.transformer = AutoModel(config)
        
        # Define the custom regression head
        transformer_output_dim = self.transformer.config.hidden_size
        self.head = nn.Sequential(
            nn.Linear(transformer_output_dim + num_extra_features, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, 1) # Output is a single regression value
        )

    def forward(self, input_ids, attention_mask, features, labels=None):
        # Pass inputs through the transformer body
        transformer_outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)
        
        # Use the [CLS] token's embedding as the text representation
        cls_embedding = transformer_outputs.last_hidden_state[:, 0, :]
        
        # Concatenate text embedding with numerical features
        combined_features = torch.cat([cls_embedding, features], dim=1)
        
        # Get the final prediction (logit) from the custom head
        logits = self.head(combined_features).squeeze(-1)

        loss = None
        if labels is not None:
            loss_fn = nn.MSELoss()
            loss = loss_fn(logits, labels.float())
            return (loss, logits)
        
        return logits