# 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

from xformers._deprecation_warning import deprecated_function

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__()
        deprecated_function(self)

    @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)
