import svetoviz_webgpu as sv

# 1. Define Parameters
vocab_size = 100
embedding_dim = 16

# 2. Define EmbeddingBag
# Efficiently computes sum/mean/max of indices in a "bag" (sentence)
ebag = nn.EmbeddingBag(num_embeddings=vocab_size, embedding_dim=embedding_dim, mode="mean")

# 3. Flattened indices + Offsets
indices = torch.tensor([1, 4, 5, 20, 31, 44, 9], dtype=torch.long)
offsets = torch.tensor([0, 3], dtype=torch.long)

def terminal_callback(buffer, message, images, files):
    # 4. Aggregated forward pass
    output = ebag(indices, offsets)

    buffer.send_system_message(f"Input Indices: {indices.tolist()}")
    buffer.send_system_message(f"Bag Offsets: {offsets.tolist()}")
    buffer.send_system_message(f"Output Tensor Shape: {list(output.shape)}")

# 5. Start the interactive session
sv.pytorch_web(module=ebag, terminal_callback=terminal_callback)