from typing import Dict
import torch
from typing import List


def unpad_text_field_tensors(
    text_field_tensors: Dict[str, Dict[str, torch.Tensor]],
    padded_idx: int = 0
) -> List[torch.Tensor]:
    """
    Unpad text field tensors and return a list which lenth is the batch size of text_field_tensor.

    Example:
        Input:
            text_field_tensors = {"tokens": {"tokens": tensor with shape: (B, S)}}
            B = num_of_batch, S = padded sequence
        Output:
            List[tensor with shape (M) * B]
            B = num_of_batch, M = unpadded sequence lenth for different sentence
    """
    text_tensor_list = []
    target_text_field_tensor = text_field_tensors["tokens"]["token_ids"]

    for sent in target_text_field_tensor:
        # sent_len = len(sent) - (sent == padded_idx).sum()
        text_tensor_list.append(sent.clone())

    if len(text_tensor_list) != 1:
        raise ValueError("Augmented but with non-valid batch_size")
    else:
        return text_tensor_list