# embeddingbag_torch.py
import torch
import torch.nn as nn
import torch.nn.functional as F


BATCH_SIZE = 1024
EMB_DIM = 256 
VOCAB_SIZE = 50000

class Model(nn.Module):
    def __init__(self, embedding_weight): 
        super().__init__()
        self.criterion = nn.EmbeddingBag(
            VOCAB_SIZE, EMB_DIM, mode='mean', sparse=False, _weight=embedding_weight
        )


    def forward(self, input: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
        return self.criterion(input, offsets)

def get_inputs():
    AVG_LEN = 10
    TOTAL_INDICES = BATCH_SIZE * AVG_LEN
    indices = torch.randint(0, VOCAB_SIZE, (TOTAL_INDICES,), dtype=torch.long)
    offsets = [0]
    current_offset = 0
    for _ in range(BATCH_SIZE - 1):
        seq_len = torch.randint(1, AVG_LEN * 2, (1,)).item()
        current_offset += seq_len
        offsets.append(current_offset)
    
    offsets = torch.tensor(offsets, dtype=torch.long)

    if offsets[-1].item() > indices.shape[0]:
        indices = indices[:offsets[-1].item()]
        
    return [indices.cuda(), offsets.cuda()]

def get_init_inputs():
    weight = torch.randn(VOCAB_SIZE, EMB_DIM, dtype=torch.float32)
    return [weight]