import svetoviz_webgpu as sv

# 1. Define parameters
vocab_size = 50
embedding_dim = 16

# 2. Define Embedding Layer
# A lookup table mapping discrete indices to continuous vectors
embedding_layer = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)

# 3. Create discrete input (Indices)
indices = torch.tensor([[10, 22, 5, 49, 0, 12, 31, 7]], dtype=torch.long)

def terminal_callback(buffer, message, images, files):
    # 4. Forward pass: O(1) lookup per index
    output = embedding_layer(indices)

    buffer.send_system_message(f"Input Token Sequence: {indices.tolist()[0]}")
    buffer.send_system_message(f"Lookup Table: {vocab_size} tokens × {embedding_dim} dimensions")
    buffer.send_system_message(f"Output Tensor Shape: {list(output.shape)}")

# 5. Start the interactive session
Svetoviz.pytorch_web(module=embedding_layer, terminal_callback=terminal_callback)