#!/usr/bin/env python3
|
# -*- coding: utf-8 -*-
|
|
# Copyright 2019 Shigeki Karita
|
# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
|
"""Unility functions for Transformer."""
|
|
import torch
|
from funasr.modules.nets_utils import pad_list
|
|
|
def add_sos_eos(ys_pad, sos, eos, ignore_id):
|
"""Add <sos> and <eos> labels.
|
|
:param torch.Tensor ys_pad: batch of padded target sequences (B, Lmax)
|
:param int sos: index of <sos>
|
:param int eos: index of <eos>
|
:param int ignore_id: index of padding
|
:return: padded tensor (B, Lmax)
|
:rtype: torch.Tensor
|
:return: padded tensor (B, Lmax)
|
:rtype: torch.Tensor
|
"""
|
|
_sos = ys_pad.new([sos])
|
_eos = ys_pad.new([eos])
|
ys = [y[y != ignore_id] for y in ys_pad] # parse padded ys
|
ys_in = [torch.cat([_sos, y], dim=0) for y in ys]
|
ys_out = [torch.cat([y, _eos], dim=0) for y in ys]
|
return pad_list(ys_in, eos), pad_list(ys_out, ignore_id)
|
