36 lines
972 B
Python
36 lines
972 B
Python
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
|
#
|
|
# This source code is licensed under the BSD license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
|
|
from abc import ABCMeta, abstractmethod
|
|
from dataclasses import asdict, dataclass
|
|
from typing import Type, TypeVar
|
|
|
|
import torch.nn as nn
|
|
|
|
Self = TypeVar("Self", bound="PositionEmbedding")
|
|
|
|
|
|
@dataclass
|
|
class PositionEmbeddingConfig:
|
|
name: str
|
|
dim_model: int
|
|
seq_len: int
|
|
|
|
|
|
class PositionEmbedding(nn.Module, metaclass=ABCMeta):
|
|
@abstractmethod
|
|
def __init__(self, *args, **kwargs) -> None:
|
|
super().__init__()
|
|
|
|
@classmethod
|
|
def from_config(cls: Type[Self], config: PositionEmbeddingConfig) -> Self:
|
|
# Generate the class inputs from the config
|
|
fields = asdict(config)
|
|
|
|
# Skip all Nones so that default values are used
|
|
fields = {k: v for k, v in fields.items() if v is not None}
|
|
return cls(**fields)
|