Function get_beam_width_array

examples/utils.py:258–289

(bwa: str = None)

def get_beam_width_array(bwa: str = None):
    bwa = ast.literal_eval(bwa)  # Short for "beam_width_array"
    if isinstance(bwa, str):
        bwa = ast.literal_eval(bwa)  # parse again for string

    def parse_one_bwa(row):
        assert isinstance(row, list), f"Beam width array must be a list."
        assert len(
            row
        ) <= 8, "Length of beam width array must not be greater than 8 now."
        assert all([isinstance(beam, int) for beam in row
                    ]), "Numbers in beam width array must be integer."
        bwa_tensor = torch.zeros([8], dtype=torch.int32)
        for j in range(len(row)):
            bwa_tensor[j] = row[j]
        bwa_tensor[len(row):] = row[-1]
        return bwa_tensor, max(row)

    if isinstance(bwa, list):  # Only one BWA
        bwa_tensor, max_beam_width = parse_one_bwa(bwa)
    elif isinstance(bwa, tuple):  # BWA for respective requests
        bwa_tensor_list = []
        max_beam_width = 0
        for row in bwa:
            bwa_tensor, beam_width = parse_one_bwa(row)
            bwa_tensor_list.append(bwa_tensor)
            max_beam_width = max(max_beam_width, beam_width)
        bwa_tensor = torch.stack(bwa_tensor_list, dim=0)
    else:
        raise ValueError(f"Invalid beam width array: {bwa}")

    return bwa_tensor.tolist(), max_beam_width


def add_common_args(parser):