|
|
| import math |
|
|
| import torch |
| import torch.utils.checkpoint |
| import torch.nn.functional as F |
| from transformers import PreTrainedModel |
|
|
| from transformers.modeling_outputs import (ModelOutput,) |
|
|
|
|
| class CounterModel(PreTrainedModel): |
|
|
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.config = config |
| self.weight = config.weight |
| self.bias = config.bias |
| self.linear = torch.nn.Linear(1,1) |
| def forward(self, x,**kwargs): |
| x = self.weight * x + self.bias |
| logits = self.linear(x) |
| return logits |
| def add(self): |
| return self.weight + self.bias |
