# Copyright (c) 2022 Agora
# Licensed under The MIT License [see LICENSE for details]

import bitsandbytes as bnb
from zeta.nn.embeddings.base import BaseEmbedding

class BnBEmbedding(BaseEmbedding):
    def forward(
        self, 
        num_tokens: int, 
        dim: int, 
        padding_idx
    ) -> bnb.nn.modules:
        embedding = bnb.nn.modules.Embedding(num_tokens, dim, padding_idx)

        return embedding
    