File size: 2,738 Bytes
e903a9a | 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 | # Define some utils for constructing Metis.
from transformers import GenerationMixin
class TrajectoryGenerationMixin(GenerationMixin):
def reset(self):
raise NotImplementedError
def commit(self):
raise NotImplementedError
def step_generate(self, input_ids, **kwargs):
outputs = self.generate(
input_ids=input_ids,
return_dict_in_generate=True,
output_hidden_states=False,
**kwargs
)
trajectory_ids = outputs.sequences
# Re-forward the trajectory to get the final hidden states.
final_outputs = self.model.forward(
input_ids=trajectory_ids,
output_hidden_states=True,
use_cache=False
)
self.commit(final_outputs)
return trajectory_ids
import torch
import torch.nn as nn
import importlib
def create_metis_decoder_layer(config, raw_decoder):
module = importlib.import_module(f"{__package__}.{config.backbone_meta['backbone_type']}_wrapper")
decoder_layer_class = getattr(module, '%sDecoderLayerForMetis' % config.backbone_meta['backbone_type'])
return decoder_layer_class(config, raw_decoder)
def create_metis_causallm(config):
module = importlib.import_module(f"{__package__}.{config.backbone_meta['backbone_type']}_wrapper")
model_class = getattr(module, '%sCausalLMForMetis' % config.backbone_meta['backbone_type'])
return model_class(config)
class DecoderLayerWrapperForMetis(nn.Module):
def __init__(self, config, raw_decoder):
super().__init__()
self.config = config
self._raw_decoder_ref = [raw_decoder]
@property
def raw_decoder(self):
return self._raw_decoder_ref[0]
def before_mixin(self, **kwargs):
raise NotImplementedError
def after_mixin(self, memory_carrier, cache_dict, **kwargs) -> torch.Tensor:
raise NotImplementedError
class CausalLMWrapperForMetis(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
def register_metis_blocks(self, metis_blocks):
# Store as a plain list so PyTorch does NOT register these as submodules
# of this wrapper. The MetisModel already owns the ModuleList; registering
# it here too would create duplicate state-dict keys (shared-tensor error).
self._metis_blocks_ref = list(metis_blocks)
def get_decoder_layer_by_id(self, layer_id: int):
raise NotImplementedError
def forward_with_memory(self, **kwargs):
raise NotImplementedError
if False: # pragma: no cover - dependency marker for HF dynamic modules
from .Qwen3_5_wrapper import Qwen3_5CausalLMForMetis
|