| from dataclasses import dataclass |
|
|
| import torch |
| from torch import nn |
| from transformers import DistilBertConfig, DistilBertModel, PretrainedConfig, PreTrainedModel |
| from transformers.utils import ModelOutput |
|
|
|
|
| DEFAULT_DECISIONS = [ |
| {'name': 'task_complete', 'prompt': 'Has the task been successfully completed?', 'options': ['false', 'true']}, |
| {'name': 'should_retry', 'prompt': 'Should the same failed action be retried without changing it?', 'options': ['false', 'true']}, |
| {'name': 'failure_source', 'prompt': 'What is the source of the original reproduced failure?', |
| 'options': ['code', 'environment', 'test', 'dependency', 'unknown']}, |
| {'name': 'next_action', 'prompt': 'What action should the agent take next under the diagnostic policy?', |
| 'options': ['inspect_code', 'rerun_tests', 'inspect_environment', 'revert_change', 'escalate']}, |
| {'name': 'evidence_sufficient', 'prompt': 'Does exactly one controlled intervention identify the original failure source?', |
| 'options': ['false', 'true']}, |
| ] |
|
|
|
|
| class SystemOneConfig(PretrainedConfig): |
| model_type = 'system_one' |
|
|
| def __init__(self, encoder_config=None, max_length=256, head_depth=2, |
| representation_dropout=0.1, head_dropout=0.1, |
| decisions=None, temperatures=None, calibration=None, **kwargs): |
| kwargs.setdefault('tie_word_embeddings', False) |
| super().__init__(**kwargs) |
| self.encoder_config = encoder_config or DistilBertConfig().to_dict() |
| self.max_length = max_length |
| self.head_depth = head_depth |
| self.representation_dropout = representation_dropout |
| self.head_dropout = head_dropout |
| self.decisions = decisions or DEFAULT_DECISIONS |
| self.num_labels_per_head = [len(item['options']) for item in self.decisions] |
| self.temperatures = temperatures |
| self.calibration = calibration |
| if self.head_depth != 2: |
| raise ValueError('This release class supports the published depth-2 checkpoint only') |
| if len(self.decisions) != 5: |
| raise ValueError('Expected exactly five fixed decisions') |
| if temperatures is not None and (len(temperatures) != 5 or any(value <= 0 for value in temperatures)): |
| raise ValueError('Expected five positive temperatures') |
|
|
|
|
| @dataclass |
| class SystemOneOutput(ModelOutput): |
| logits: tuple[torch.Tensor, ...] | None = None |
| hidden_states: tuple[torch.Tensor, ...] | None = None |
| attentions: tuple[torch.Tensor, ...] | None = None |
|
|
|
|
| class SystemOneModel(PreTrainedModel): |
| config_class = SystemOneConfig |
| base_model_prefix = 'encoder' |
| main_input_name = 'input_ids' |
| _supports_sdpa = True |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| encoder_config = DistilBertConfig.from_dict(config.encoder_config) |
| if config._attn_implementation is not None: |
| encoder_config._attn_implementation = config._attn_implementation |
| self.encoder = DistilBertModel(encoder_config) |
| self.dropout = nn.Dropout(config.representation_dropout) |
| self.heads = nn.ModuleList([ |
| nn.Sequential( |
| nn.Linear(encoder_config.dim, encoder_config.dim), |
| nn.GELU(), |
| nn.Dropout(config.head_dropout), |
| nn.Linear(encoder_config.dim, outputs), |
| ) |
| for outputs in config.num_labels_per_head |
| ]) |
| self.post_init() |
|
|
| def forward(self, input_ids=None, attention_mask=None, head_mask=None, inputs_embeds=None, |
| output_attentions=None, output_hidden_states=None, return_dict=None): |
| return_dict = self.config.use_return_dict if return_dict is None else return_dict |
| outputs = self.encoder( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| head_mask=head_mask, |
| inputs_embeds=inputs_embeds, |
| output_attentions=output_attentions, |
| output_hidden_states=output_hidden_states, |
| return_dict=True, |
| ) |
| representation = outputs.last_hidden_state[:, 0] |
| logits = tuple(head(self.dropout(representation)) for head in self.heads) |
| if not return_dict: |
| return logits, outputs.hidden_states, outputs.attentions |
| return SystemOneOutput(logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions) |
|
|
