wehub-resource-sync
97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from fairseq.data import Dictionary
|
|
from fairseq.models import (
|
|
FairseqDecoder,
|
|
FairseqLanguageModel,
|
|
register_model,
|
|
register_model_architecture,
|
|
)
|
|
|
|
|
|
@register_model("dummy_model")
|
|
class DummyModel(FairseqLanguageModel):
|
|
def __init__(self, args, encoder):
|
|
super().__init__(encoder)
|
|
self.args = args
|
|
|
|
@staticmethod
|
|
def add_args(parser):
|
|
parser.add_argument("--num-layers", type=int, default=24)
|
|
parser.add_argument("--embed-dim", type=int, default=1024)
|
|
|
|
@classmethod
|
|
def build_model(cls, args, task):
|
|
encoder = DummyEncoder(
|
|
num_embed=len(task.target_dictionary),
|
|
embed_dim=args.embed_dim,
|
|
num_layers=args.num_layers,
|
|
)
|
|
return cls(args, encoder)
|
|
|
|
def forward(self, src_tokens, masked_tokens=None, **kwargs):
|
|
return self.decoder(src_tokens, masked_tokens=masked_tokens)
|
|
|
|
|
|
class DummyEncoder(FairseqDecoder):
|
|
def __init__(self, num_embed=50000, embed_dim=1024, num_layers=24):
|
|
super().__init__(Dictionary())
|
|
self.embed = nn.Embedding(
|
|
num_embeddings=num_embed, embedding_dim=embed_dim, padding_idx=0
|
|
)
|
|
self.layers_a = nn.ModuleList(
|
|
[
|
|
nn.Sequential(
|
|
nn.LayerNorm(embed_dim),
|
|
nn.Linear(embed_dim, 3 * embed_dim), # q, k, v input projection
|
|
nn.Linear(3 * embed_dim, embed_dim), # skip self-attention
|
|
nn.Linear(embed_dim, embed_dim), # output projection
|
|
nn.Dropout(),
|
|
)
|
|
for i in range(num_layers)
|
|
]
|
|
)
|
|
self.layers_b = nn.ModuleList(
|
|
[
|
|
nn.Sequential(
|
|
nn.LayerNorm(embed_dim),
|
|
nn.Linear(embed_dim, 4 * embed_dim), # FFN
|
|
nn.ReLU(),
|
|
nn.Linear(4 * embed_dim, embed_dim), # FFN
|
|
nn.Dropout(0.1),
|
|
)
|
|
for i in range(num_layers)
|
|
]
|
|
)
|
|
self.out_proj = nn.Linear(embed_dim, num_embed)
|
|
|
|
def forward(self, tokens, masked_tokens=None):
|
|
x = self.embed(tokens)
|
|
for layer_a, layer_b in zip(self.layers_a, self.layers_b):
|
|
x = x + layer_a(x)
|
|
x = x + layer_b(x)
|
|
x = self.out_proj(x)
|
|
if masked_tokens is not None:
|
|
x = x[masked_tokens]
|
|
return (x,)
|
|
|
|
def max_positions(self):
|
|
return 1024
|
|
|
|
def get_normalized_probs(self, net_output, log_probs, sample=None):
|
|
logits = net_output[0].float()
|
|
if log_probs:
|
|
return F.log_softmax(logits, dim=-1)
|
|
else:
|
|
return F.softmax(logits, dim=-1)
|
|
|
|
|
|
@register_model_architecture("dummy_model", "dummy_model")
|
|
def base_architecture(args):
|
|
pass
|