| 24 | |
| 25 | @dataclass |
| 26 | class TransformerBlockMetadata: |
| 27 | return_hidden_states_index: int = None |
| 28 | return_encoder_hidden_states_index: int = None |
| 29 | hidden_states_argument_name: str = "hidden_states" |
| 30 | |
| 31 | _cls: Type = None |
| 32 | _cached_parameter_indices: dict[str, int] = None |
| 33 | |
| 34 | def _get_parameter_from_args_kwargs(self, identifier: str, args=(), kwargs=None): |
| 35 | kwargs = kwargs or {} |
| 36 | if identifier in kwargs: |
| 37 | return kwargs[identifier] |
| 38 | if self._cached_parameter_indices is not None: |
| 39 | return args[self._cached_parameter_indices[identifier]] |
| 40 | if self._cls is None: |
| 41 | raise ValueError("Model class is not set for metadata.") |
| 42 | parameters = list(inspect.signature(self._cls.forward).parameters.keys()) |
| 43 | parameters = parameters[1:] # skip `self` |
| 44 | self._cached_parameter_indices = {param: i for i, param in enumerate(parameters)} |
| 45 | if identifier not in self._cached_parameter_indices: |
| 46 | raise ValueError(f"Parameter '{identifier}' not found in function signature but was requested.") |
| 47 | index = self._cached_parameter_indices[identifier] |
| 48 | if index >= len(args): |
| 49 | raise ValueError(f"Expected {index} arguments but got {len(args)}.") |
| 50 | return args[index] |
| 51 | |
| 52 | |
| 53 | class AttentionProcessorRegistry: |
no outgoing calls
searching dependent graphs…