|
from math import sqrt,log |
|
import sys |
|
|
|
|
|
|
|
import torch |
|
import torch.nn as nn |
|
from torch.nn.functional import softmax,relu,linear |
|
from common import PositionalEncoding |
|
from hopfield import HopfieldLayer, HopfieldMHA, HopfieldReLU, HopfieldSoftmax |
|
|
|
from torch.cuda.amp import autocast |
|
import yaml |
|
from transformers import PreTrainedModel, PretrainedConfig |
|
from transformers.modeling_outputs import MaskedLMOutput, BaseModelOutput |
|
|
|
class BertEnergyConfig(PretrainedConfig): |
|
|
|
model_type = "bert_energy" |
|
|
|
def __init__(self, config=None, path=None, vocabulary_size=50, num_layers=12, num_heads=12, forward_memories=2048, embedding_dim=768, activation="relu",positional=True, bias=True, tie_weights=True, alpha=1.0, |
|
beta=1., layer_norm=1e-05, dropout=0.0, block_size=512, share_layers=False, compile=False, pad_idx=None, **kwargs): |
|
|
|
self.vocabulary_size = vocabulary_size |
|
self.num_layers = num_layers |
|
self.num_heads = num_heads |
|
self.activation = activation |
|
self.positional = positional |
|
self.tie_weights = tie_weights |
|
self.bias = bias |
|
self.forward_memories = forward_memories |
|
self.embedding_dim = embedding_dim |
|
self.share_layers = share_layers |
|
self.alpha = alpha |
|
self.beta = beta |
|
self.layer_norm = float(layer_norm) |
|
self.dropout = dropout |
|
self.block_size = block_size |
|
self.compile = compile |
|
self.pad_idx = pad_idx |
|
|
|
if config is not None: |
|
for key,value in config.to_dict(): |
|
if key.lower() in self.__dict__.keys(): |
|
print(key, file=sys.stderr) |
|
setattr(self,key.lower(),value) |
|
|
|
elif path is not None: |
|
if path.endswith(".yaml"): |
|
with open(path) as istream: |
|
config = yaml.safe_load(istream) |
|
for key,value in config.items(): |
|
print(key) |
|
if key.lower() in self.__dict__.keys(): |
|
setattr(self,key.lower(),value) |
|
else: |
|
raise NotImplementedError |
|
super().__init__(**kwargs) |
|
|
