| """Hub modeling file for the assistant-memory sentence classifier. |
| |
| This file is copied verbatim into the published model repository as |
| ``modeling_stencil.py`` and loaded through ``trust_remote_code``. It depends only on |
| ``torch`` and ``transformers``. The registered classifier scores the raw ``[CLS]`` |
| hidden state (no BERT pooler) with a linear head that also receives a one-hot speaker |
| role; the role part of that head is stored here as a per-role logit bias so that |
| ``logits = W_text @ cls + b + role_bias[role]`` reproduces the original numbers. |
| |
| Inputs are sentence PAIRS, exactly as registered: ``text`` = the preceding context or |
| ``"(no context)"``, ``text_pair`` = ``"[<role>] <sentence>"``. Labels: rule, fact, none. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
| from transformers import BertConfig, BertModel, BertPreTrainedModel |
| from transformers.modeling_outputs import SequenceClassifierOutput |
|
|
|
|
| class StencilSentenceClassifierConfig(BertConfig): |
| model_type = "stencil_sentence_classifier" |
|
|
| def __init__(self, roles=("user", "assistant"), default_role="user", **kwargs): |
| super().__init__(**kwargs) |
| self.roles = list(roles) |
| self.default_role = default_role |
|
|
|
|
| class StencilSentenceClassifier(BertPreTrainedModel): |
| config_class = StencilSentenceClassifierConfig |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.bert = BertModel(config, add_pooling_layer=False) |
| self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels) |
| self.role_bias = torch.nn.Parameter( |
| torch.zeros(len(config.roles), config.num_labels) |
| ) |
| self.post_init() |
|
|
| def forward( |
| self, |
| input_ids=None, |
| attention_mask=None, |
| token_type_ids=None, |
| role_ids=None, |
| labels=None, |
| **kwargs, |
| ): |
| hidden = self.bert( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| token_type_ids=token_type_ids, |
| ).last_hidden_state[:, 0] |
| logits = self.classifier(hidden) |
| if role_ids is None: |
| role_ids = torch.full( |
| (logits.shape[0],), |
| self.config.roles.index(self.config.default_role), |
| dtype=torch.long, |
| device=logits.device, |
| ) |
| logits = logits + self.role_bias[role_ids] |
| loss = None |
| if labels is not None: |
| loss = torch.nn.functional.cross_entropy(logits, labels) |
| return SequenceClassifierOutput(loss=loss, logits=logits) |
|
|
