from typing import List
from typing import Dict
import torch


def batch_to_word_ids(batch: List[List[str]], vocab2id: Dict[str, int]) -> torch.Tensor:
    """
    For word embedding.

    1. UNK will be mapped to 3 since we assume the vocab starts with `<S>, </S>, <UNK>`.

    Return tensor of shape `(batch_size, max_timesteps)`.
    """
    max_timesteps = max(len(row) for row in batch)

    rows = []
    for words in batch:
        row = [vocab2id.get(word, 3) for word in words]
        row.extend([0] * (max_timesteps - len(row)))
        rows.append(row)

    return torch.LongTensor(rows)