"""
A BERT model that predicts how complete an answer is.

The head is a single linear layer on the CLS token, with no pooler in between.
That is not one of the standard transformers heads, so the class lives here and
you load it with trust_remote_code=True.
"""

import torch
from torch import nn
from transformers import BertModel, BertPreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput


class BertCompletenessRegressor(BertPreTrainedModel):
    """Predicts a completeness score. Higher means more complete."""

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = 1
        self.bert = BertModel(config, add_pooling_layer=False)
        self.regressor = nn.Linear(config.hidden_size, 1)
        self.post_init()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
                position_ids=None, head_mask=None, inputs_embeds=None,
                labels=None, output_attentions=None, output_hidden_states=None,
                return_dict=None):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.bert(
            input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids,
            position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds,
            output_attentions=output_attentions, output_hidden_states=output_hidden_states,
            return_dict=True,
        )
        cls = outputs.last_hidden_state[:, 0]   # CLS token, no pooler
        logits = self.regressor(cls)

        loss = None
        if labels is not None:
            loss = nn.functional.mse_loss(logits.squeeze(-1), labels.float())

        if not return_dict:
            return (loss, logits) if loss is not None else (logits,)

        return SequenceClassifierOutput(
            loss=loss, logits=logits,
            hidden_states=outputs.hidden_states, attentions=outputs.attentions,
        )


def build_input(question, answer):
    """Format a question and answer the way the model was trained.

    The stray "f" before "Answer:" is a typo in the original training code.
    It has to stay, otherwise the input does not match what the model saw.
    """
    return f'Question: {question}\n\nfAnswer: {answer}\n\nHow complete is this answer?'