File size: 4,395 Bytes
0b4d63e 3df6849 0b4d63e 3df6849 0b4d63e 3df6849 0b4d63e 3df6849 0b4d63e 3df6849 0b4d63e 3df6849 0b4d63e 3df6849 0b4d63e 3df6849 0b4d63e 3df6849 0b4d63e 3df6849 0b4d63e 3df6849 0b4d63e 3df6849 0b4d63e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | 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)
|