File size: 3,958 Bytes
e34b94f |
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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
import torch
import torch.nn as nn
from transformers import (
PreTrainedModel,
AutoModelForCausalLM,
AutoTokenizer,
Cache
)
from transformers.utils import TransformersKwargs
from transformers.processing_utils import Unpack
from typing import Optional
from abc import ABC, abstractmethod
from peft import PeftConfig, get_peft_model
class Trigger(torch.nn.Module, ABC):
def __init__(self):
pass
@abstractmethod
def forward(self, **kwargs) -> bool:
...
class NanoTrigger(torch.nn.Module):
def __init__(self):
super().__init__()
self.register_buffer("_device", torch.tensor(0.0))
@property
def device(self):
return self._device.device
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> bool:
# This "nano trigger" always predicts insertion.
# It outputs logits where the probability of insertion (index=1) is set to 1.0
# for every token position in the batch.
batch_size, seq_len = input_ids.shape
logits = torch.zeros(batch_size, seq_len, 2, device=input_ids.device)
logits[..., 1] = 1.0
return logits
class MemGenTrigger(torch.nn.Module):
"""
Trigger module for the MemGen Model.
- Input: the weaver receives `inputs_embeds` from the reasoner model's current decoding sequence.
- Output: the weaver produces a sequence of hidden states with length K,
which are concatenated to the original `inputs_embeds` to alter the reasoner's decoding path.
"""
def __init__(
self,
pretrained_model_name_or_path: str,
peft_config: Optional[PeftConfig] = None
):
super().__init__()
# base model
self.model = AutoModelForCausalLM.from_pretrained(
pretrained_model_name_or_path, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2"
)
self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)
self.model = self._postprocess(self.model)
if peft_config is not None:
self.model = get_peft_model(self.model, peft_config)
self.config = self.model.config
@property
def device(self):
return self.model.device
def _postprocess(self, model: PreTrainedModel):
for parameter in model.parameters():
parameter.requires_grad = True
# Replace lm_head with a binary classification head
hidden_size = model.config.hidden_size
classification_head = nn.Linear(hidden_size, 2)
model.lm_head = classification_head
# Ensure the new head parameters are trainable
for param in model.lm_head.parameters():
param.requires_grad = True
return model
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
attention_mask: Optional[torch.Tensor] = None,
**kwargs: Unpack[TransformersKwargs],
) -> torch.Tensor:
"""Trigger decision mechanism for sequence generation.
The trigger determines its decision based on the already generated `input_ids`.
It is influenced by the data distribution but is independent of the weaver module.
Args:
input_ids (Optional[torch.LongTensor]): Token IDs of the generated sequence.
attention_mask (Optional[torch.Tensor]): Attention mask to avoid attending to padding tokens. Defaults to None.
**kwargs: Additional keyword arguments passed to the underlying model.
Returns:
torch.Tensor: Logits tensor of shape `(batch_size, seq_len, num_classes)`,
representing the trigger's decision probabilities.
"""
return self.model(
input_ids=input_ids,
attention_mask=attention_mask,
**kwargs
).logits |