Duplicate from nguyenvulebinh/spoken-norm-taggen

Co-authored-by: Binh Nguyen <nguyenvulebinh@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,27 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,38 @@
+---
+title: Spoken Norm
+emoji: 📊
+colorFrom: gray
+colorTo: yellow
+sdk: gradio
+app_file: app.py
+pinned: false
+duplicated_from: nguyenvulebinh/spoken-norm-taggen
+---
+
+# Configuration
+
+`title`: _string_  
+Display title for the Space
+
+`emoji`: _string_  
+Space emoji (emoji-only character allowed)
+
+`colorFrom`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+
+`colorTo`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+
+`sdk`: _string_  
+Can be either `gradio` or `streamlit`
+
+`sdk_version` : _string_  
+Only applicable for `streamlit` SDK.  
+See [doc](https://hf.co/docs/hub/spaces) for more info on supported versions.
+
+`app_file`: _string_  
+Path to your main application file (which contains either `gradio` or `streamlit` Python code).  
+Path is relative to the root of the repository.
+
+`pinned`: _boolean_  
+Whether the Space stays on top of your list.

app.py

@@ -0,0 +1,25 @@
+import gradio as gr
+from infer import infer
+
+
+
+def format_text(text_input, list_bias_input):
+    print('{}\n{}\n\n'.format(text_input, list_bias_input))
+    bias_list = list_bias_input.strip().split('\n')
+    norm_result = infer([text_input], bias_list)
+    return norm_result[0]
+
+
+title = "Transformation spoken text to written text"
+
+iface = gr.Interface(format_text,
+                     [
+                         gr.inputs.Textbox(
+                             lines=1,
+                             default="ngày hai tám tháng tư cô vít bùng phát ở xì cút len chiếm tám mươi phần trăm là biến chủng đen ta và bê ta và ô mi cờ ron"),
+                         gr.inputs.Textbox(
+                             lines=5, default='covid\ndelta\nbeta\nomicron | ô mi cờ ron\nscotland | sờ cốt lờn | xì cút len'),
+                     ],
+                     outputs="text",
+                     title=title)
+iface.launch()

attentions.py

@@ -0,0 +1,466 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+import numpy as np
+from typing import Optional, Tuple
+
+
+class ScaledDotProductAttention(nn.Module):
+    """
+    Scaled Dot-Product Attention proposed in "Attention Is All You Need"
+    Compute the dot products of the query with all keys, divide each by sqrt(dim),
+    and apply a softmax function to obtain the weights on the values
+
+    Args: dim, mask
+        dim (int): dimention of attention
+        mask (torch.Tensor): tensor containing indices to be masked
+
+    Inputs: query, key, value, mask
+        - **query** (batch, q_len, d_model): tensor containing projection vector for decoder.
+        - **key** (batch, k_len, d_model): tensor containing projection vector for encoder.
+        - **value** (batch, v_len, d_model): tensor containing features of the encoded input sequence.
+        - **mask** (-): tensor containing indices to be masked
+
+    Returns: context, attn
+        - **context**: tensor containing the context vector from attention mechanism.
+        - **attn**: tensor containing the attention (alignment) from the encoder outputs.
+    """
+
+    def __init__(self, dim: int):
+        super(ScaledDotProductAttention, self).__init__()
+        self.sqrt_dim = np.sqrt(dim)
+
+    def forward(self, query: Tensor, key: Tensor, value: Tensor, mask: Optional[Tensor] = None) -> Tuple[
+        Tensor, Tensor]:
+        score = torch.bmm(query, key.transpose(1, 2)) / self.sqrt_dim
+
+        if mask is not None:
+            score.masked_fill_(mask.view(score.size()), -float('Inf'))
+
+        attn = F.softmax(score, -1)
+        context = torch.bmm(attn, value)
+        return context, attn
+
+
+class DotProductAttention(nn.Module):
+    """
+    Compute the dot products of the query with all values and apply a softmax function to obtain the weights on the values
+    """
+
+    def __init__(self, hidden_dim):
+        super(DotProductAttention, self).__init__()
+        self.normalize = nn.LayerNorm(hidden_dim)
+
+    def forward(self, query: Tensor, value: Tensor) -> Tuple[Tensor, Tensor]:
+        batch_size, hidden_dim, input_size = query.size(0), query.size(2), value.size(1)
+
+        score = torch.bmm(query, value.transpose(1, 2))
+        attn = F.softmax(score.view(-1, input_size), dim=1).view(batch_size, -1, input_size)
+        context = torch.bmm(attn, value)
+
+        return context, attn
+
+
+class AdditiveAttention(nn.Module):
+    """
+     Applies a additive attention (bahdanau) mechanism on the output features from the decoder.
+     Additive attention proposed in "Neural Machine Translation by Jointly Learning to Align and Translate" paper.
+
+     Args:
+         hidden_dim (int): dimesion of hidden state vector
+
+     Inputs: query, value
+         - **query** (batch_size, q_len, hidden_dim): tensor containing the output features from the decoder.
+         - **value** (batch_size, v_len, hidden_dim): tensor containing features of the encoded input sequence.
+
+     Returns: context, attn
+         - **context**: tensor containing the context vector from attention mechanism.
+         - **attn**: tensor containing the alignment from the encoder outputs.
+
+     Reference:
+         - **Neural Machine Translation by Jointly Learning to Align and Translate**: https://arxiv.org/abs/1409.0473
+    """
+
+    def __init__(self, hidden_dim: int) -> None:
+        super(AdditiveAttention, self).__init__()
+        self.query_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.key_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.bias = nn.Parameter(torch.rand(hidden_dim).uniform_(-0.1, 0.1))
+        self.score_proj = nn.Linear(hidden_dim, 1)
+
+    def forward(self, query: Tensor, key: Tensor, value: Tensor) -> Tuple[Tensor, Tensor]:
+        score = self.score_proj(torch.tanh(self.key_proj(key) + self.query_proj(query) + self.bias)).squeeze(-1)
+        attn = F.softmax(score, dim=-1)
+        context = torch.bmm(attn.unsqueeze(1), value)
+        return context, attn
+
+
+class LocationAwareAttention(nn.Module):
+    """
+    Applies a location-aware attention mechanism on the output features from the decoder.
+    Location-aware attention proposed in "Attention-Based Models for Speech Recognition" paper.
+    The location-aware attention mechanism is performing well in speech recognition tasks.
+    We refer to implementation of ClovaCall Attention style.
+
+    Args:
+        hidden_dim (int): dimesion of hidden state vector
+        smoothing (bool): flag indication whether to use smoothing or not.
+
+    Inputs: query, value, last_attn, smoothing
+        - **query** (batch, q_len, hidden_dim): tensor containing the output features from the decoder.
+        - **value** (batch, v_len, hidden_dim): tensor containing features of the encoded input sequence.
+        - **last_attn** (batch_size * num_heads, v_len): tensor containing previous timestep`s attention (alignment)
+
+    Returns: output, attn
+        - **output** (batch, output_len, dimensions): tensor containing the feature from encoder outputs
+        - **attn** (batch * num_heads, v_len): tensor containing the attention (alignment) from the encoder outputs.
+
+    Reference:
+        - **Attention-Based Models for Speech Recognition**: https://arxiv.org/abs/1506.07503
+        - **ClovaCall**: https://github.com/clovaai/ClovaCall/blob/master/las.pytorch/models/attention.py
+    """
+
+    def __init__(self, hidden_dim: int, smoothing: bool = True) -> None:
+        super(LocationAwareAttention, self).__init__()
+        self.hidden_dim = hidden_dim
+        self.conv1d = nn.Conv1d(in_channels=1, out_channels=hidden_dim, kernel_size=3, padding=1)
+        self.query_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.value_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.score_proj = nn.Linear(hidden_dim, 1, bias=True)
+        self.bias = nn.Parameter(torch.rand(hidden_dim).uniform_(-0.1, 0.1))
+        self.smoothing = smoothing
+
+    def forward(self, query: Tensor, value: Tensor, last_attn: Tensor) -> Tuple[Tensor, Tensor]:
+        batch_size, hidden_dim, seq_len = query.size(0), query.size(2), value.size(1)
+
+        # Initialize previous attention (alignment) to zeros
+        if last_attn is None:
+            last_attn = value.new_zeros(batch_size, seq_len)
+
+        conv_attn = torch.transpose(self.conv1d(last_attn.unsqueeze(1)), 1, 2)
+        score = self.score_proj(torch.tanh(
+            self.query_proj(query.reshape(-1, hidden_dim)).view(batch_size, -1, hidden_dim)
+            + self.value_proj(value.reshape(-1, hidden_dim)).view(batch_size, -1, hidden_dim)
+            + conv_attn
+            + self.bias
+        )).squeeze(dim=-1)
+
+        if self.smoothing:
+            score = torch.sigmoid(score)
+            attn = torch.div(score, score.sum(dim=-1).unsqueeze(dim=-1))
+        else:
+            attn = F.softmax(score, dim=-1)
+
+        context = torch.bmm(attn.unsqueeze(dim=1), value).squeeze(dim=1)  # Bx1xT X BxTxD => Bx1xD => BxD
+
+        return context, attn
+
+
+class MultiHeadLocationAwareAttention(nn.Module):
+    """
+    Applies a multi-headed location-aware attention mechanism on the output features from the decoder.
+    Location-aware attention proposed in "Attention-Based Models for Speech Recognition" paper.
+    The location-aware attention mechanism is performing well in speech recognition tasks.
+    In the above paper applied a signle head, but we applied multi head concept.
+
+    Args:
+        hidden_dim (int): The number of expected features in the output
+        num_heads (int): The number of heads. (default: )
+        conv_out_channel (int): The number of out channel in convolution
+
+    Inputs: query, value, prev_attn
+        - **query** (batch, q_len, hidden_dim): tensor containing the output features from the decoder.
+        - **value** (batch, v_len, hidden_dim): tensor containing features of the encoded input sequence.
+        - **prev_attn** (batch_size * num_heads, v_len): tensor containing previous timestep`s attention (alignment)
+
+    Returns: output, attn
+        - **output** (batch, output_len, dimensions): tensor containing the feature from encoder outputs
+        - **attn** (batch * num_heads, v_len): tensor containing the attention (alignment) from the encoder outputs.
+
+    Reference:
+        - **Attention Is All You Need**: https://arxiv.org/abs/1706.03762
+        - **Attention-Based Models for Speech Recognition**: https://arxiv.org/abs/1506.07503
+    """
+
+    def __init__(self, hidden_dim: int, num_heads: int = 8, conv_out_channel: int = 10) -> None:
+        super(MultiHeadLocationAwareAttention, self).__init__()
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.dim = int(hidden_dim / num_heads)
+        self.conv1d = nn.Conv1d(num_heads, conv_out_channel, kernel_size=3, padding=1)
+        self.loc_proj = nn.Linear(conv_out_channel, self.dim, bias=False)
+        self.query_proj = nn.Linear(hidden_dim, self.dim * num_heads, bias=False)
+        self.value_proj = nn.Linear(hidden_dim, self.dim * num_heads, bias=False)
+        self.score_proj = nn.Linear(self.dim, 1, bias=True)
+        self.bias = nn.Parameter(torch.rand(self.dim).uniform_(-0.1, 0.1))
+
+    def forward(self, query: Tensor, value: Tensor, last_attn: Tensor) -> Tuple[Tensor, Tensor]:
+        batch_size, seq_len = value.size(0), value.size(1)
+
+        if last_attn is None:
+            last_attn = value.new_zeros(batch_size, self.num_heads, seq_len)
+
+        loc_energy = torch.tanh(self.loc_proj(self.conv1d(last_attn).transpose(1, 2)))
+        loc_energy = loc_energy.unsqueeze(1).repeat(1, self.num_heads, 1, 1).view(-1, seq_len, self.dim)
+
+        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.dim).permute(0, 2, 1, 3)
+        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.dim).permute(0, 2, 1, 3)
+        query = query.contiguous().view(-1, 1, self.dim)
+        value = value.contiguous().view(-1, seq_len, self.dim)
+
+        score = self.score_proj(torch.tanh(value + query + loc_energy + self.bias)).squeeze(2)
+        attn = F.softmax(score, dim=1)
+
+        value = value.view(batch_size, seq_len, self.num_heads, self.dim).permute(0, 2, 1, 3)
+        value = value.contiguous().view(-1, seq_len, self.dim)
+
+        context = torch.bmm(attn.unsqueeze(1), value).view(batch_size, -1, self.num_heads * self.dim)
+        attn = attn.view(batch_size, self.num_heads, -1)
+
+        return context, attn
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-Head Attention proposed in "Attention Is All You Need"
+    Instead of performing a single attention function with d_model-dimensional keys, values, and queries,
+    project the queries, keys and values h times with different, learned linear projections to d_head dimensions.
+    These are concatenated and once again projected, resulting in the final values.
+    Multi-head attention allows the model to jointly attend to information from different representation
+    subspaces at different positions.
+
+    MultiHead(Q, K, V) = Concat(head_1, ..., head_h) · W_o
+        where head_i = Attention(Q · W_q, K · W_k, V · W_v)
+
+    Args:
+        d_model (int): The dimension of keys / values / quries (default: 512)
+        num_heads (int): The number of attention heads. (default: 8)
+
+    Inputs: query, key, value, mask
+        - **query** (batch, q_len, d_model): In transformer, three different ways:
+            Case 1: come from previoys decoder layer
+            Case 2: come from the input embedding
+            Case 3: come from the output embedding (masked)
+
+        - **key** (batch, k_len, d_model): In transformer, three different ways:
+            Case 1: come from the output of the encoder
+            Case 2: come from the input embeddings
+            Case 3: come from the output embedding (masked)
+
+        - **value** (batch, v_len, d_model): In transformer, three different ways:
+            Case 1: come from the output of the encoder
+            Case 2: come from the input embeddings
+            Case 3: come from the output embedding (masked)
+
+        - **mask** (-): tensor containing indices to be masked
+
+    Returns: output, attn
+        - **output** (batch, output_len, dimensions): tensor containing the attended output features.
+        - **attn** (batch * num_heads, v_len): tensor containing the attention (alignment) from the encoder outputs.
+    """
+
+    def __init__(self, d_model: int = 512, num_heads: int = 8):
+        super(MultiHeadAttention, self).__init__()
+
+        assert d_model % num_heads == 0, "d_model % num_heads should be zero."
+
+        self.d_head = int(d_model / num_heads)
+        self.num_heads = num_heads
+        self.scaled_dot_attn = ScaledDotProductAttention(self.d_head)
+        self.query_proj = nn.Linear(d_model, self.d_head * num_heads)
+        self.key_proj = nn.Linear(d_model, self.d_head * num_heads)
+        self.value_proj = nn.Linear(d_model, self.d_head * num_heads)
+
+    def forward(
+            self,
+            query: Tensor,
+            key: Tensor,
+            value: Tensor,
+            mask: Optional[Tensor] = None
+    ) -> Tuple[Tensor, Tensor]:
+        batch_size = value.size(0)
+
+        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)  # BxQ_LENxNxD
+        key = self.key_proj(key).view(batch_size, -1, self.num_heads, self.d_head)  # BxK_LENxNxD
+        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.d_head)  # BxV_LENxNxD
+
+        query = query.permute(2, 0, 1, 3).contiguous().view(batch_size * self.num_heads, -1, self.d_head)  # BNxQ_LENxD
+        key = key.permute(2, 0, 1, 3).contiguous().view(batch_size * self.num_heads, -1, self.d_head)  # BNxK_LENxD
+        value = value.permute(2, 0, 1, 3).contiguous().view(batch_size * self.num_heads, -1, self.d_head)  # BNxV_LENxD
+
+        if mask is not None:
+            mask = mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1)  # BxNxQ_LENxK_LEN
+
+        context, attn = self.scaled_dot_attn(query, key, value, mask)
+
+        context = context.view(self.num_heads, batch_size, -1, self.d_head)
+        context = context.permute(1, 2, 0, 3).contiguous().view(batch_size, -1, self.num_heads * self.d_head)  # BxTxND
+
+        return context, attn
+
+
+class RelativeMultiHeadAttention(nn.Module):
+    """
+    Multi-head attention with relative positional encoding.
+    This concept was proposed in the "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context"
+
+    Args:
+        d_model (int): The dimension of model
+        num_heads (int): The number of attention heads.
+        dropout_p (float): probability of dropout
+
+    Inputs: query, key, value, pos_embedding, mask
+        - **query** (batch, time, dim): Tensor containing query vector
+        - **key** (batch, time, dim): Tensor containing key vector
+        - **value** (batch, time, dim): Tensor containing value vector
+        - **pos_embedding** (batch, time, dim): Positional embedding tensor
+        - **mask** (batch, 1, time2) or (batch, time1, time2): Tensor containing indices to be masked
+
+    Returns:
+        - **outputs**: Tensor produces by relative multi head attention module.
+    """
+
+    def __init__(
+            self,
+            d_model: int = 512,
+            num_heads: int = 16,
+            dropout_p: float = 0.1,
+    ):
+        super(RelativeMultiHeadAttention, self).__init__()
+        assert d_model % num_heads == 0, "d_model % num_heads should be zero."
+        self.d_model = d_model
+        self.d_head = int(d_model / num_heads)
+        self.num_heads = num_heads
+        self.sqrt_dim = math.sqrt(d_model)
+
+        self.query_proj = nn.Linear(d_model, d_model)
+        self.key_proj = nn.Linear(d_model, d_model)
+        self.value_proj = nn.Linear(d_model, d_model)
+        self.pos_proj = nn.Linear(d_model, d_model, bias=False)
+
+        self.dropout = nn.Dropout(p=dropout_p)
+        self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))
+        self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))
+        torch.nn.init.xavier_uniform_(self.u_bias)
+        torch.nn.init.xavier_uniform_(self.v_bias)
+
+        self.out_proj = nn.Linear(d_model, d_model)
+
+    def forward(
+            self,
+            query: Tensor,
+            key: Tensor,
+            value: Tensor,
+            pos_embedding: Tensor,
+            mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        batch_size = value.size(0)
+
+        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)
+        key = self.key_proj(key).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)
+        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)
+        pos_embedding = self.pos_proj(pos_embedding).view(batch_size, -1, self.num_heads, self.d_head)
+
+        content_score = torch.matmul((query + self.u_bias).transpose(1, 2), key.transpose(2, 3))
+        pos_score = torch.matmul((query + self.v_bias).transpose(1, 2), pos_embedding.permute(0, 2, 3, 1))
+        pos_score = self._compute_relative_positional_encoding(pos_score)
+
+        score = (content_score + pos_score) / self.sqrt_dim
+
+        if mask is not None:
+            mask = mask.unsqueeze(1)
+            score.masked_fill_(mask, -1e9)
+
+        attn = F.softmax(score, -1)
+        attn = self.dropout(attn)
+
+        context = torch.matmul(attn, value).transpose(1, 2)
+        context = context.contiguous().view(batch_size, -1, self.d_model)
+
+        return self.out_proj(context)
+
+    def _compute_relative_positional_encoding(self, pos_score: Tensor) -> Tensor:
+        batch_size, num_heads, seq_length1, seq_length2 = pos_score.size()
+        zeros = pos_score.new_zeros(batch_size, num_heads, seq_length1, 1)
+        padded_pos_score = torch.cat([zeros, pos_score], dim=-1)
+
+        padded_pos_score = padded_pos_score.view(batch_size, num_heads, seq_length2 + 1, seq_length1)
+        pos_score = padded_pos_score[:, :, 1:].view_as(pos_score)
+
+        return pos_score
+
+
+class CustomizingAttention(nn.Module):
+    r"""
+    Customizing Attention
+
+    Applies a multi-head + location-aware attention mechanism on the output features from the decoder.
+    Multi-head attention proposed in "Attention Is All You Need" paper.
+    Location-aware attention proposed in "Attention-Based Models for Speech Recognition" paper.
+    I combined these two attention mechanisms as custom.
+
+    Args:
+        hidden_dim (int): The number of expected features in the output
+        num_heads (int): The number of heads. (default: )
+        conv_out_channel (int): The dimension of convolution
+
+    Inputs: query, value, last_attn
+        - **query** (batch, q_len, hidden_dim): tensor containing the output features from the decoder.
+        - **value** (batch, v_len, hidden_dim): tensor containing features of the encoded input sequence.
+        - **last_attn** (batch_size * num_heads, v_len): tensor containing previous timestep`s alignment
+
+    Returns: output, attn
+        - **output** (batch, output_len, dimensions): tensor containing the attended output features from the decoder.
+        - **attn** (batch * num_heads, v_len): tensor containing the alignment from the encoder outputs.
+
+    Reference:
+        - **Attention Is All You Need**: https://arxiv.org/abs/1706.03762
+        - **Attention-Based Models for Speech Recognition**: https://arxiv.org/abs/1506.07503
+    """
+
+    def __init__(self, hidden_dim: int, num_heads: int = 4, conv_out_channel: int = 10) -> None:
+        super(CustomizingAttention, self).__init__()
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.dim = int(hidden_dim / num_heads)
+        self.scaled_dot_attn = ScaledDotProductAttention(self.dim)
+        self.conv1d = nn.Conv1d(1, conv_out_channel, kernel_size=3, padding=1)
+        self.query_proj = nn.Linear(hidden_dim, self.dim * num_heads, bias=True)
+        self.value_proj = nn.Linear(hidden_dim, self.dim * num_heads, bias=False)
+        self.loc_proj = nn.Linear(conv_out_channel, self.dim, bias=False)
+        self.bias = nn.Parameter(torch.rand(self.dim * num_heads).uniform_(-0.1, 0.1))
+
+    def forward(self, query: Tensor, value: Tensor, last_attn: Tensor) -> Tuple[Tensor, Tensor]:
+        batch_size, q_len, v_len = value.size(0), query.size(1), value.size(1)
+
+        if last_attn is None:
+            last_attn = value.new_zeros(batch_size * self.num_heads, v_len)
+
+        loc_energy = self.get_loc_energy(last_attn, batch_size, v_len)  # get location energy
+
+        query = self.query_proj(query).view(batch_size, q_len, self.num_heads * self.dim)
+        value = self.value_proj(value).view(batch_size, v_len, self.num_heads * self.dim) + loc_energy + self.bias
+
+        query = query.view(batch_size, q_len, self.num_heads, self.dim).permute(2, 0, 1, 3)
+        value = value.view(batch_size, v_len, self.num_heads, self.dim).permute(2, 0, 1, 3)
+        query = query.contiguous().view(-1, q_len, self.dim)
+        value = value.contiguous().view(-1, v_len, self.dim)
+
+        context, attn = self.scaled_dot_attn(query, value)
+        attn = attn.squeeze()
+
+        context = context.view(self.num_heads, batch_size, q_len, self.dim).permute(1, 2, 0, 3)
+        context = context.contiguous().view(batch_size, q_len, -1)
+
+        return context, attn
+
+    def get_loc_energy(self, last_attn: Tensor, batch_size: int, v_len: int) -> Tensor:
+        conv_feat = self.conv1d(last_attn.unsqueeze(1))
+        conv_feat = conv_feat.view(batch_size, self.num_heads, -1, v_len).permute(0, 1, 3, 2)
+
+        loc_energy = self.loc_proj(conv_feat).view(batch_size, self.num_heads, v_len, self.dim)
+        loc_energy = loc_energy.permute(0, 2, 1, 3).reshape(batch_size, v_len, self.num_heads * self.dim)
+
+        return loc_energy

data_handling.py

@@ -0,0 +1,336 @@
+import datasets
+import model_handling
+from transformers import PreTrainedTokenizerBase
+from typing import Optional, Union, Any
+from transformers.file_utils import PaddingStrategy
+import re
+import os
+from tqdm import tqdm
+# import time
+import json
+import random
+import regtag
+from dataclasses import dataclass
+import validators
+
+import utils
+
+regexp = re.compile(r"\d{4}[\-/]\d{2}[\-/]\d{2}t\d{2}:\d{2}:\d{2}")
+target_bias_words = set(regtag.get_general_en_word())
+tokenizer = None
+
+
+def get_bias_words():
+    regtag.augment.get_random_oov()
+    return list(regtag.augment.oov_dict.keys())
+
+
+def check_common_phrase(word):
+    if validators.email(word.replace(' @', '@')):
+        return True
+    if validators.domain(word):
+        return True
+    if validators.url(word):
+        return True
+    if word in regtag.get_general_en_word():
+        return True
+    return False
+
+
+@dataclass
+class DataCollatorForNormSeq2Seq:
+    tokenizer: PreTrainedTokenizerBase
+    model: Optional[Any] = None
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    label_pad_token_id: int = -100
+    return_tensors: str = "pt"
+
+    def bias_phrases_extractor(self, features, max_bias_per_sample=30):
+        # src_ids, src_length, tgt_ids, tgt_length
+        phrase_candidate = []
+        sample_output_words = []
+        bias_labels = []
+
+        for sample in features:
+            words = []
+            for idx, (src_word_len, tgt_word_len) in enumerate(zip(sample['inputs_length'], sample['outputs_length'])):
+                src_start_idx = sum(sample['inputs_length'][:idx])
+                tgt_start_idx = sum(sample['outputs_length'][:idx])
+                word_input = self.tokenizer.decode(sample['input_ids'][src_start_idx: src_start_idx + src_word_len])
+                word_output = self.tokenizer.decode(sample['outputs'][tgt_start_idx: tgt_start_idx + tgt_word_len])
+                words.append(word_output)
+                if word_input != word_output and not any(map(str.isdigit, word_output)):
+                    phrase_candidate.append(word_output)
+            sample_output_words.append(words)
+
+        phrase_candidate = list(set(phrase_candidate))
+        phrase_candidate_revised = []
+        phrase_candidate_common = []
+        raw_phrase_candidate = []
+        for item in phrase_candidate:
+            raw_item = self.tokenizer.sp_model.DecodePieces(item.split())
+            if check_common_phrase(raw_item):
+                phrase_candidate_common.append(raw_item)
+            else:
+                phrase_candidate_revised.append(item)
+            raw_phrase_candidate.append(raw_item)
+
+        remain_phrase = max(0, max_bias_per_sample * len(features) - len(phrase_candidate_revised))
+
+        if remain_phrase > 0:
+            words_candidate = list(
+                set(get_bias_words()) - set(raw_phrase_candidate))
+            random.shuffle(words_candidate)
+            phrase_candidate_revised += [' '.join(self.tokenizer.sp_model.EncodeAsPieces(item)[:5]) for item in
+                                         words_candidate[:remain_phrase]]
+
+        for i in range(len(features)):
+            sample_bias_lables = []
+            for w_idx, w in enumerate(sample_output_words[i]):
+                try:
+                    sample_bias_lables.extend(
+                        [phrase_candidate_revised.index(w) + 1] * features[i]['outputs_length'][w_idx])
+                except:
+                    # random ignore 0 label
+                    if random.random() < 0.5:
+                        sample_bias_lables.extend([0] * features[i]['outputs_length'][w_idx])
+                    else:
+                        sample_bias_lables.extend([self.label_pad_token_id] * features[i]['outputs_length'][w_idx])
+            bias_labels.append(sample_bias_lables)
+            assert len(sample_bias_lables) == len(features[i]['outputs']), "{} vs {}".format(sample_bias_lables,
+                                                                                             features[i]['outputs'])
+
+        # phrase_candidate_ids = [self.tokenizer.encode(item) for item in phrase_candidate]
+        phrase_candidate_ids = [self.tokenizer.encode(self.tokenizer.sp_model.DecodePieces(item.split())) for item in
+                                phrase_candidate_revised]
+        phrase_candidate_mask = [[self.tokenizer.pad_token_id] * len(item) for item in phrase_candidate_ids]
+
+        return phrase_candidate_ids, phrase_candidate_mask, bias_labels
+        # pass
+
+    def encode_list_string(self, list_text):
+        text_tokenized = self.tokenizer(list_text)
+        return self.tokenizer.pad(
+            text_tokenized,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors='pt',
+        )
+
+    def __call__(self, features, return_tensors=None):
+        # start_time = time.time()
+        batch_src, batch_tgt = [], []
+        for item in features:
+            src_spans, tgt_spans = utils.make_spoken(item['text'])
+            batch_src.append(src_spans)
+            batch_tgt.append(tgt_spans)
+        # print("Make src-tgt {}s".format(time.time() - start_time))
+        # start_time = time.time()
+
+        features = preprocess_function({"src": batch_src, "tgt": batch_tgt})
+
+
+        # print("Make feature {}s".format(time.time() - start_time))
+        # start_time = time.time()
+
+        phrase_candidate_ids, phrase_candidate_mask, samples_bias_labels = self.bias_phrases_extractor(features)
+        # print("Make bias {}s".format(time.time() - start_time))
+        # start_time = time.time()
+
+        if return_tensors is None:
+            return_tensors = self.return_tensors
+        labels = [feature["outputs"] for feature in features] if "outputs" in features[0].keys() else None
+        spoken_labels = [feature["spoken_label"] for feature in features] if "spoken_label" in features[0].keys() else None
+        spoken_idx = [feature["src_spoken_idx"] for feature in features] if "src_spoken_idx" in features[0].keys() else None
+
+        word_src_lengths = [feature["inputs_length"] for feature in features] if "inputs_length" in features[0].keys() else None
+        word_tgt_lengths = [feature["outputs_length"] for feature in features] if "outputs_length" in features[0].keys() else None
+        # We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the
+        # same length to return tensors.
+        if labels is not None:
+            max_label_length = max(len(l) for l in labels)
+            max_src_length = max(len(l) for l in spoken_labels)
+            max_spoken_idx_length = max(len(l) for l in spoken_idx)
+            max_word_src_length = max(len(l) for l in word_src_lengths)
+            max_word_tgt_length = max(len(l) for l in word_tgt_lengths)
+
+            padding_side = self.tokenizer.padding_side
+            for feature, bias_labels in zip(features, samples_bias_labels):
+                remainder = [self.label_pad_token_id] * (max_label_length - len(feature["outputs"]))
+                remainder_word_tgt_length = [0] * (max_word_tgt_length - len(feature["outputs_length"]))
+                remainder_spoken = [self.label_pad_token_id] * (max_src_length - len(feature["spoken_label"]))
+                remainder_spoken_idx = [self.label_pad_token_id] * (max_spoken_idx_length - len(feature["src_spoken_idx"]))
+                remainder_word_src_length = [0] * (max_word_src_length - len(feature["inputs_length"]))
+
+                feature["labels"] = (
+                    feature["outputs"] + [
+                        self.tokenizer.eos_token_id] + remainder if padding_side == "right" else remainder + feature[
+                        "outputs"] + [self.tokenizer.eos_token_id]
+                )
+                feature["labels_bias"] = (
+                    bias_labels + [0] + remainder if padding_side == "right" else remainder + bias_labels + [0]
+                )
+
+                feature["spoken_label"] = [self.label_pad_token_id] + feature["spoken_label"] + [self.label_pad_token_id]
+                feature["spoken_label"] = feature["spoken_label"] + remainder_spoken if padding_side == "right" else remainder_spoken + feature["spoken_label"]
+                feature["src_spoken_idx"] = feature["src_spoken_idx"] + remainder_spoken_idx
+
+                feature['inputs_length'] = [1] + feature['inputs_length'] + [1]
+                feature['outputs_length'] = feature['outputs_length'] + [1]
+
+                feature["inputs_length"] = feature["inputs_length"] + remainder_word_src_length
+                feature["outputs_length"] = feature["outputs_length"] + remainder_word_tgt_length
+
+
+        features_inputs = [{
+            "input_ids": [self.tokenizer.bos_token_id] + item["input_ids"] + [self.tokenizer.eos_token_id],
+            "attention_mask": [self.tokenizer.pad_token_id] + item["attention_mask"] + [self.tokenizer.pad_token_id]
+        } for item in features]
+        features_inputs = self.tokenizer.pad(
+            features_inputs,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors=return_tensors,
+        )
+
+        bias_phrases_inputs = [{
+            "input_ids": ids,
+            "attention_mask": mask
+        } for ids, mask in zip(phrase_candidate_ids, phrase_candidate_mask)]
+        bias_phrases_inputs = self.tokenizer.pad(
+            bias_phrases_inputs,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors=return_tensors,
+        )
+
+        outputs = self.tokenizer.pad({"input_ids": [feature["labels"] for feature in features]},
+                                     return_tensors=return_tensors)['input_ids']
+        outputs_bias = self.tokenizer.pad({"input_ids": [feature["labels_bias"] for feature in features]},
+                                          return_tensors=return_tensors)['input_ids']
+        spoken_label = self.tokenizer.pad({"input_ids": [feature["spoken_label"] for feature in features]},
+                                          return_tensors=return_tensors)['input_ids']
+        spoken_idx = self.tokenizer.pad({"input_ids": [feature["src_spoken_idx"] for feature in features]},
+                                        return_tensors=return_tensors)['input_ids'] + 1  # 1 for bos token
+        word_src_lengths = self.tokenizer.pad({"input_ids": [feature["inputs_length"] for feature in features]},
+                                              return_tensors=return_tensors)['input_ids']
+        word_tgt_lengths = self.tokenizer.pad({"input_ids": [feature["outputs_length"] for feature in features]},
+                                              return_tensors=return_tensors)['input_ids']
+
+        features = {
+            "input_ids": features_inputs["input_ids"],
+            "spoken_label": spoken_label,
+            "spoken_idx": spoken_idx,
+            "word_src_lengths": word_src_lengths,
+            "word_tgt_lengths": word_tgt_lengths,
+            "attention_mask": features_inputs["attention_mask"],
+            "bias_input_ids": bias_phrases_inputs["input_ids"],
+            "bias_attention_mask": bias_phrases_inputs["attention_mask"],
+            "labels": outputs,
+            "labels_bias": outputs_bias
+        }
+
+        # print("Make batch {}s".format(time.time() - start_time))
+        # start_time = time.time()
+
+        # prepare decoder_input_ids
+        if self.model is not None and hasattr(self.model, "prepare_decoder_input_ids_from_labels"):
+            decoder_input_ids = self.model.prepare_decoder_input_ids_from_labels(labels=features["labels"])
+            features["decoder_input_ids"] = decoder_input_ids
+
+        return features
+
+
+# data init
+def init_data(train_corpus_path='./data-bin/raw/train_raw.txt',
+              test_corpus_path='./data-bin/raw/valid_raw.txt'):
+    dataset_oov = datasets.load_dataset('text', data_files={"train": train_corpus_path,
+                                                            "test": test_corpus_path})
+
+    print(dataset_oov)
+    return dataset_oov
+
+
+def preprocess_function(batch):
+
+    global tokenizer
+    if tokenizer is None:
+        tokenizer = model_handling.init_tokenizer()
+
+    features = []
+    for src_words, tgt_words in zip(batch["src"], batch["tgt"]):
+        src_ids, pad_ids, src_lengths, tgt_ids, tgt_lengths = [], [], [], [], []
+        src_spoken_label = []  # 0: "O", 1: "B", 2: "I"
+
+        src_spoken_idx = []
+        tgt_spoken_ids = []
+
+        for idx, (src, tgt) in enumerate(zip(src_words, tgt_words)):
+            is_remain = False
+            if src == tgt:
+                is_remain = True
+
+            src_tokenized = tokenizer(src)
+            if len(src_tokenized['input_ids']) < 3:
+                continue
+            # hardcode fix tokenizer email
+            if validators.email(tgt):
+                tgt_tokenized = tokenizer(tgt.replace('@', ' @'))
+            else:
+                tgt_tokenized = tokenizer(tgt)
+            if len(tgt_tokenized['input_ids']) < 3:
+                continue
+            src_ids.extend(src_tokenized["input_ids"][1:-1])
+            if is_remain:
+                src_spoken_label.extend([0 if random.random() < 0.5 else -100 for _ in range(len(src_tokenized["input_ids"][1:-1]))])
+                if random.random() < 0.1:
+                    # Random pick normal word for spoken norm
+                    src_spoken_idx.append(idx)
+                    tgt_spoken_ids.append(tgt_tokenized["input_ids"][1:-1])
+            else:
+                src_spoken_label.extend([1] + [2] * (len(src_tokenized["input_ids"][1:-1]) - 1))
+                src_spoken_idx.append(idx)
+                tgt_spoken_ids.append(tgt_tokenized["input_ids"][1:-1])
+
+            pad_ids.extend(src_tokenized["attention_mask"][1:-1])
+            src_lengths.append(len(src_tokenized["input_ids"]) - 2)
+            tgt_ids.extend(tgt_tokenized["input_ids"][1:-1])
+            tgt_lengths.append(len(tgt_tokenized["input_ids"]) - 2)
+            if len(src_ids) > 70 or len(tgt_ids) > 70:
+                # print("Ignore sample")
+                break
+                
+        if len(src_ids) < 1 or len(tgt_ids) < 1:
+            continue
+        # else:
+            # print("ignore")
+            
+        features.append({
+            "input_ids": src_ids,
+            "attention_mask": pad_ids,
+            "spoken_label": src_spoken_label,
+            "inputs_length": src_lengths,
+            "outputs": tgt_ids,
+            "outputs_length": tgt_lengths,
+            "src_spoken_idx": src_spoken_idx,
+            "tgt_spoken_ids": tgt_spoken_ids
+        })
+
+    return features
+
+
+if __name__ == "__main__":
+    split_datasets = init_data()
+
+    model, model_tokenizer = model_handling.init_model()
+    data_collator = DataCollatorForNormSeq2Seq(model_tokenizer, model=model)
+
+    # start = time.time()
+    batch = data_collator([split_datasets["train"][i] for i in [random.randint(0, 900) for _ in range(0, 12)]])
+    print(batch)
+    # print("{}s".format(time.time() - start))

infer.py

@@ -0,0 +1,374 @@
+#!/usr/bin/env python
+# coding: utf-8
+import torch
+import model_handling
+from data_handling import DataCollatorForNormSeq2Seq
+from model_handling import EncoderDecoderSpokenNorm
+import os
+import random
+import data_handling
+from transformers.generation_logits_process import LogitsProcessorList
+from transformers.generation_stopping_criteria import StoppingCriteriaList
+from transformers.generation_beam_search import BeamSearchScorer
+from dataclasses import dataclass
+from transformers.file_utils import ModelOutput
+import utils
+
+# os.environ["CUDA_VISIBLE_DEVICES"] = "4"
+
+use_gpu = False
+if use_gpu:
+    if not torch.cuda.is_available():
+        use_gpu = False
+tokenizer = model_handling.init_tokenizer()
+model = EncoderDecoderSpokenNorm.from_pretrained('nguyenvulebinh/spoken-norm-taggen-v2').eval()
+data_collator = DataCollatorForNormSeq2Seq(tokenizer)
+if use_gpu:
+    model = model.cuda()
+
+
+def make_batch_input(text_input_list):
+    batch_src_ids, batch_src_lengths = [], []
+    for text_input in text_input_list:
+        src_ids, src_lengths = [], []
+        for src in text_input.split():
+            src_tokenized = tokenizer(src)
+            ids = src_tokenized["input_ids"][1:-1]
+            src_ids.extend(ids)
+            src_lengths.append(len(ids))
+        src_ids = torch.tensor([0] + src_ids + [2])
+        src_lengths = torch.tensor([1] + src_lengths + [1]) + 1
+        batch_src_ids.append(src_ids)
+        batch_src_lengths.append(src_lengths)
+        assert sum(src_lengths - 1) == len(src_ids), "{} vs {}".format(sum(src_lengths), len(src_ids))
+    input_tokenized = tokenizer.pad({"input_ids": batch_src_ids}, padding=True)
+    input_word_length = tokenizer.pad({"input_ids": batch_src_lengths}, padding=True)["input_ids"] - 1
+    return input_tokenized['input_ids'], input_tokenized['attention_mask'], input_word_length
+
+
+def make_batch_bias_list(bias_list):
+    if len(bias_list) > 0:
+        bias = data_collator.encode_list_string(bias_list)
+        bias_input_ids = bias['input_ids']
+        bias_attention_mask = bias['attention_mask']
+    else:
+        bias_input_ids = None
+        bias_attention_mask = None
+
+    return bias_input_ids, bias_attention_mask
+
+
+def build_spoken_pronounce_mapping(bias_list):
+    list_pronounce = []
+    mapping = dict({})
+    for item in bias_list:
+        pronounces = item.split(' | ')[1:]
+        pronounces = [tokenizer(item)['input_ids'][1:-1] for item in pronounces]
+        list_pronounce.extend(pronounces)    
+    subword_ids = list(set([item for sublist in list_pronounce for item in sublist]))
+    mapping = {item: [] for item in subword_ids}
+    for item in list_pronounce:
+        for wid in subword_ids:
+            if wid in item:
+                mapping[wid].append(item)
+    return mapping
+
+def find_pivot(seq, subseq):
+    n = len(seq)
+    m = len(subseq)
+    result = []
+    for i in range(n - m + 1):
+        if seq[i] == subseq[0] and seq[i:i + m] == subseq:
+            result.append(i)
+    return result
+
+def revise_spoken_tagging(list_tags, list_words, pronounce_mapping):
+    if len(pronounce_mapping) == 0:
+        return list_tags
+    result = []
+    for tags_tensor, sen in zip(list_tags, list_words):
+        tags = tags_tensor.detach().numpy().tolist()
+        sen = sen.detach().numpy().tolist()
+        candidate_pronounce = dict({})
+        for idx in range(len(tags)):
+            if tags[idx] != 0 and sen[idx] in pronounce_mapping:
+                for pronounce in pronounce_mapping[sen[idx]]:
+                    pronounce_word = str(pronounce)
+                    start_find_idx = max(0, idx - len(pronounce))
+                    end_find_idx = idx + len(pronounce)
+                    find_idx = find_pivot(sen[start_find_idx: end_find_idx], pronounce)
+                    if len(find_idx) > 0:
+                        find_idx = [item + start_find_idx for item in find_idx]
+                        for map_idx in find_idx:
+                            if candidate_pronounce.get(map_idx, None) is None:
+                                candidate_pronounce[map_idx] = len(pronounce)
+                            else:
+                                candidate_pronounce[map_idx] = max(candidate_pronounce[map_idx], len(pronounce))
+        for idx, len_word in candidate_pronounce.items():
+            tags_tensor[idx] = 1
+            for i in range(1, len_word):
+                tags_tensor[idx + i] = 2
+        result.append(tags_tensor)
+    return result
+
+
+def make_spoken_feature(input_features, text_input_list, pronounce_mapping=dict({})):
+    features = {
+        "input_ids": input_features[0],
+        "word_src_lengths": input_features[2],
+        "attention_mask": input_features[1],
+        # "bias_input_ids": bias_features[0],
+        # "bias_attention_mask": bias_features[1],
+        "bias_input_ids": None,
+        "bias_attention_mask": None,
+    }
+    if use_gpu:
+        for key in features.keys():
+            if features[key] is not None:
+                features[key] = features[key].cuda()
+        
+    encoder_output = model.get_encoder()(**features)
+    spoken_tagging_output = torch.argmax(encoder_output[0].spoken_tagging_output, dim=-1)
+    spoken_tagging_output = revise_spoken_tagging(spoken_tagging_output, features['input_ids'], pronounce_mapping)
+    
+    # print(spoken_tagging_output)
+    # print(features['input_ids'])
+    word_src_lengths = features['word_src_lengths']
+    encoder_features = encoder_output[0][0]
+    list_spoken_features = []
+    list_pre_norm = []
+    for tagging_sample, sample_word_length, text_input_features, sample_text in zip(spoken_tagging_output, word_src_lengths, encoder_features, text_input_list):
+        spoken_feature_idx = []
+        sample_words = ['<s>'] + sample_text.split() + ['</s>']
+        norm_words = []
+        spoken_phrase = []
+        spoken_features = []
+        if tagging_sample.sum() == 0:
+            list_pre_norm.append(sample_words)
+            continue
+        for idx, word_length in enumerate(sample_word_length):
+            if word_length > 0:
+                start = sample_word_length[:idx].sum()
+                end = start + word_length
+                if tagging_sample[start: end].sum() > 0 and sample_words[idx] not in ['<s>', '</s>']:
+                    # Word has start tag
+                    if (tagging_sample[start: end] == 1).sum():
+                        if len(spoken_phrase) > 0:
+                            norm_words.append('<mask>[{}]({})'.format(len(list_spoken_features), ' '.join(spoken_phrase)))
+                            spoken_phrase = []
+                            list_spoken_features.append(torch.cat(spoken_features))
+                            spoken_features = []
+                    spoken_phrase.append(sample_words[idx]) 
+                    spoken_features.append(text_input_features[start: end])
+                else:
+                    if len(spoken_phrase) > 0:
+                        norm_words.append('<mask>[{}]({})'.format(len(list_spoken_features), ' '.join(spoken_phrase)))
+                        spoken_phrase = []
+                        list_spoken_features.append(torch.cat(spoken_features))
+                        spoken_features = []
+                    norm_words.append(sample_words[idx])
+        if len(spoken_phrase) > 0:
+            norm_words.append('<mask>[{}]({})'.format(len(list_spoken_features), ' '.join(spoken_phrase)))
+            spoken_phrase = []
+            list_spoken_features.append(torch.cat(spoken_features))
+            spoken_features = []
+        list_pre_norm.append(norm_words)
+        
+        
+    list_features_mask = []
+    if len(list_spoken_features) > 0:
+        feature_pad = torch.zeros_like(list_spoken_features[0][:1, :])
+        max_length = max([len(item) for item in list_spoken_features])
+        for i in range(len(list_spoken_features)):
+            spoken_length = len(list_spoken_features[i])
+            remain_length = max_length - spoken_length
+            device = list_spoken_features[i].device
+            list_spoken_features[i] = torch.cat([list_spoken_features[i], 
+                                                 feature_pad.expand(remain_length, feature_pad.size(-1))]).unsqueeze(0)
+            list_features_mask.append(torch.cat([torch.ones(spoken_length, device=device, dtype=torch.int64),
+                                                 torch.zeros(remain_length, device=device, dtype=torch.int64)]).unsqueeze(0))
+    if len(list_spoken_features) > 0:
+        list_spoken_features = torch.cat(list_spoken_features)
+        list_features_mask = torch.cat(list_features_mask)
+    
+    return list_spoken_features, list_features_mask, list_pre_norm
+
+
+def make_bias_feature(bias_raw_features):
+    features = {
+        "bias_input_ids": bias_raw_features[0],
+        "bias_attention_mask": bias_raw_features[1]
+    }
+    if use_gpu:
+        for key in features.keys():
+            if features[key] is not None:
+                features[key] = features[key].cuda()
+    return model.forward_bias(**features)
+
+
+def decode_plain_output(decoder_output):
+    plain_output = [item.split()[1:] for item in tokenizer.batch_decode(decoder_output['sequences'], skip_special_tokens=False)]
+    scores = torch.stack(list(decoder_output['scores'])).transpose(1, 0)
+    logit_output = torch.gather(scores, -1, decoder_output['sequences'][:, 1:].unsqueeze(-1)).squeeze(-1)
+    special_tokens = list(tokenizer.special_tokens_map.values())
+    generated_output = []
+    generated_scores = []
+    # filter special tokens
+    for out_text, out_score in zip(plain_output, logit_output):
+        temp_str, tmp_score = [], []
+        for piece, score in zip(out_text, out_score):
+            if piece not in special_tokens:
+                temp_str.append(piece)
+                tmp_score.append(score)
+        if len(temp_str) > 0:
+            generated_output.append(' '.join(temp_str).replace('▁', '|').replace(' ', '').replace('|', ' ').strip())
+            generated_scores.append((sum(tmp_score)/len(tmp_score)).cpu().detach().numpy().tolist())
+        else:
+            generated_output.append("")
+            generated_scores.append(0)
+    return generated_output, generated_scores
+
+
+def generate_spoken_norm(list_spoken_features, list_features_mask, bias_features):
+    @dataclass
+    class EncoderOutputs(ModelOutput):
+        last_hidden_state: torch.FloatTensor = None
+        hidden_states: torch.FloatTensor = None
+        attentions: torch.FloatTensor = None
+
+    batch_size = list_spoken_features.size(0)
+    max_length = 50
+    device = list_spoken_features.device
+    decoder_input_ids = torch.zeros((batch_size, 1), device=device, dtype=torch.int64)
+    stopping_criteria = model._get_stopping_criteria(max_length=max_length, max_time=None,
+                                                     stopping_criteria=StoppingCriteriaList())
+    model_kwargs = {
+        "encoder_outputs": EncoderOutputs(last_hidden_state=list_spoken_features),
+        "encoder_bias_outputs": bias_features,
+        "attention_mask": list_features_mask
+    }
+    decoder_output = model.greedy_search(
+        decoder_input_ids,
+        logits_processor=LogitsProcessorList(),
+        stopping_criteria=stopping_criteria,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+        output_scores=True,
+        return_dict_in_generate=True,
+        **model_kwargs,
+    )
+    plain_output, plain_score = decode_plain_output(decoder_output)
+    # plain_output = tokenizer.batch_decode(decoder_output['sequences'], skip_special_tokens=True)
+    # # print(decoder_output)
+    # plain_output = [word.replace('▁', '|').replace(' ', '').replace('|', ' ').strip() for word in plain_output]
+    return plain_output, plain_score
+
+
+def generate_beam_spoken_norm(list_spoken_features, list_features_mask, bias_features, num_beams=3):
+    @dataclass
+    class EncoderOutputs(ModelOutput):
+        last_hidden_state: torch.FloatTensor = None
+
+    batch_size = list_spoken_features.size(0)
+    max_length = 50
+    num_return_sequences = 1
+    device = list_spoken_features.device
+    decoder_input_ids = torch.zeros((batch_size, 1), device=device, dtype=torch.int64)
+    stopping_criteria = model._get_stopping_criteria(max_length=max_length, max_time=None,
+                                                     stopping_criteria=StoppingCriteriaList())
+    model_kwargs = {
+        "encoder_outputs": EncoderOutputs(last_hidden_state=list_spoken_features),
+        "encoder_bias_outputs": bias_features,
+        "attention_mask": list_features_mask
+    }
+    beam_scorer = BeamSearchScorer(
+        batch_size=batch_size,
+        num_beams=num_beams,
+        device=device,
+        do_early_stopping=True,
+        num_beam_hyps_to_keep=num_return_sequences,
+    )
+    decoder_input_ids, model_kwargs = model._expand_inputs_for_generation(
+        decoder_input_ids, expand_size=num_beams, is_encoder_decoder=True, **model_kwargs
+    )
+
+    decoder_output = model.beam_search(
+        decoder_input_ids,
+        beam_scorer,
+        logits_processor=LogitsProcessorList(),
+        stopping_criteria=stopping_criteria,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+        output_scores=None,
+        return_dict_in_generate=True,
+        **model_kwargs,
+    )
+
+    plain_output = tokenizer.batch_decode(decoder_output['sequences'], skip_special_tokens=True)
+    plain_output = [word.replace('▁', '|').replace(' ', '').replace('|', ' ').strip() for word in plain_output]
+    return plain_output, None
+
+
+def reformat_normed_term(list_pre_norm, spoken_norm_output, spoken_norm_output_score=None, threshold=None, debug=False):
+    output = []
+    for pre_norm in list_pre_norm:
+        normed_words = []
+        # words = pre_norm.split()
+        for w in pre_norm:
+            if w.startswith('<mask>'):
+                term = w[7:].split('](')
+                # print(w)
+                # print(term)
+                term_idx = int(term[0])
+                norm_val = spoken_norm_output[term_idx]
+                norm_val_score = None if (spoken_norm_output_score is None or threshold is None) else spoken_norm_output_score[term_idx]
+                pre_norm_val = term[1][:-1]
+                if debug:
+                    if norm_val_score is not None:
+                        normed_words.append("({})({:.2f})[{}]".format(norm_val, norm_val_score, pre_norm_val))
+                    else:
+                        normed_words.append("({})[{}]".format(norm_val, pre_norm_val))
+                else:
+                    if threshold is not None and norm_val_score is not None:
+                        if norm_val_score > threshold:
+                            normed_words.append(norm_val)
+                        else:
+                            normed_words.append(pre_norm_val)
+                    else:
+                        normed_words.append(norm_val)
+            else:
+                normed_words.append(w)
+        output.append(" ".join(normed_words))
+    return output
+
+
+def infer(text_input_list, bias_list):
+    # extract bias feature
+    bias_raw_features = make_batch_bias_list(bias_list)
+    bias_features = make_bias_feature(bias_raw_features)
+    pronounce_mapping = build_spoken_pronounce_mapping(bias_list)
+
+    # Chunk split input and create feature
+    text_input_chunk_list = [utils.split_chunk_input(item, chunk_size=60, overlap=20) for item in text_input_list]
+    num_chunks = [len(i) for i in text_input_chunk_list]
+    flatten_list = [y for x in text_input_chunk_list for y in x]
+    input_raw_features = make_batch_input(flatten_list)
+
+    # Extract norm term and spoken feature
+    list_spoken_features, list_features_mask, list_pre_norm = make_spoken_feature(input_raw_features, flatten_list, pronounce_mapping)
+
+    # Merge overlap chunks
+    list_pre_norm_by_input = []
+    for idx, input_num in enumerate(num_chunks):
+        start = sum(num_chunks[:idx])
+        end = start + num_chunks[idx]
+        list_pre_norm_by_input.append(list_pre_norm[start:end])
+    text_input_list_pre_norm = [utils.merge_chunk_pre_norm(list_chunks, overlap=20, debug=False) for list_chunks in list_pre_norm_by_input]
+
+    if len(list_spoken_features) > 0:
+        spoken_norm_output, spoken_norm_score = generate_spoken_norm(list_spoken_features, list_features_mask, bias_features)
+    else:
+        spoken_norm_output, spoken_norm_score = [], None
+
+    return reformat_normed_term(text_input_list_pre_norm, spoken_norm_output, spoken_norm_score, threshold=15, debug=False)
+    

model_config_handling.py

@@ -0,0 +1,90 @@
+# coding=utf-8
+# Copyright 2020 The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers import BertConfig
+from transformers.utils import logging
+# from model_handling import DecoderSpokenNorm
+
+
+logger = logging.get_logger(__name__)
+
+
+class DecoderSpokenNormConfig(BertConfig):
+    # model_type = "decoder-spoken-norm"
+
+    def __init__(self, pad_token_id=1, bos_token_id=0, eos_token_id=2, **kwargs):
+        """Constructs RobertaConfig."""
+        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+        self.num_hidden_layers=2
+        # self.hidden_layers_from_pretrained = list(range(self.num_hidden_layers))
+        # self.hidden_layers_from_pretrained = [0, 3]
+
+        # if len(self.hidden_layers_from_pretrained) < self.num_hidden_layers:
+            # self.num_hidden_layers = len(self.hidden_layers_from_pretrained)
+
+
+class EncoderDecoderSpokenNormConfig(PretrainedConfig):
+    # model_type = "encoder-decoder-spoken-norm"
+    is_composition = True
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        assert (
+            "encoder" in kwargs and "decoder" in kwargs
+        ), "Config has to be initialized with encoder and decoder config"
+        encoder_config = kwargs.pop("encoder")
+        encoder_model_type = encoder_config.pop("model_type")
+        decoder_config = kwargs.pop("decoder")
+        decoder_model_type = decoder_config.pop("model_type")
+
+        from transformers.models.auto.configuration_auto import AutoConfig
+
+        self.encoder = AutoConfig.for_model(encoder_model_type, **encoder_config)
+        self.decoder = AutoConfig.for_model(decoder_model_type, **decoder_config)
+        self.is_encoder_decoder = True
+
+    @classmethod
+    def from_encoder_decoder_configs(
+        cls, encoder_config: PretrainedConfig, decoder_config: PretrainedConfig, **kwargs
+    ) -> PretrainedConfig:
+        r"""
+        Instantiate a :class:`~transformers.EncoderDecoderConfig` (or a derived class) from a pre-trained encoder model
+        configuration and decoder model configuration.
+
+        Returns:
+            :class:`EncoderDecoderConfig`: An instance of a configuration object
+        """
+        logger.info("Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config")
+        decoder_config.is_decoder = True
+        decoder_config.add_cross_attention = True
+
+        return cls(encoder=encoder_config.to_dict(), decoder=decoder_config.to_dict(), **kwargs)
+
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default `to_dict()` from `PretrainedConfig`.
+
+        Returns:
+            :obj:`Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+        """
+        output = copy.deepcopy(self.__dict__)
+        output["encoder"] = self.encoder.to_dict()
+        output["decoder"] = self.decoder.to_dict()
+        output["model_type"] = self.__class__.model_type
+        return output

model_handling.py

@@ -0,0 +1,763 @@
+from transformers.file_utils import cached_path, hf_bucket_url
+from importlib.machinery import SourceFileLoader
+import os
+from transformers import EncoderDecoderModel, AutoConfig, AutoModel, EncoderDecoderConfig, RobertaForCausalLM, \
+    RobertaModel
+from transformers.modeling_utils import PreTrainedModel, logging
+import torch
+from torch.nn import CrossEntropyLoss, Parameter
+from transformers.modeling_outputs import Seq2SeqLMOutput, CausalLMOutputWithCrossAttentions, \
+    ModelOutput
+from attentions import ScaledDotProductAttention, MultiHeadAttention
+from collections import namedtuple
+from typing import Dict, Any, Optional, Tuple
+from dataclasses import dataclass
+import random
+from model_config_handling import EncoderDecoderSpokenNormConfig, DecoderSpokenNormConfig, PretrainedConfig
+
+cache_dir = './cache'
+model_name = 'nguyenvulebinh/envibert'
+
+if not os.path.exists(cache_dir):
+    os.makedirs(cache_dir)
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+class SpokenNormOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    logits_spoken_tagging: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
+    cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
+    encoder_last_hidden_state: Optional[torch.FloatTensor] = None
+    encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+
+
+def collect_spoken_phrases_features(encoder_hidden_states, word_src_lengths, spoken_label):
+    list_features = []
+    list_features_mask = []
+    max_length = word_src_lengths.max()
+    feature_pad = torch.zeros_like(encoder_hidden_states[0, :1, :])
+    for hidden_state, word_length, list_idx in zip(encoder_hidden_states, word_src_lengths, spoken_label):
+        for idx in list_idx:
+            if idx > 0:
+                start = sum(word_length[:idx])
+                end = start + word_length[idx]
+                remain_length = max_length - word_length[idx]
+                list_features_mask.append(torch.cat([torch.ones_like(spoken_label[0, 0]).expand(word_length[idx]),
+                                                     torch.zeros_like(
+                                                         spoken_label[0, 0].expand(remain_length))]).unsqueeze(0))
+                spoken_phrases_feature = hidden_state[start: end]
+
+                list_features.append(torch.cat([spoken_phrases_feature,
+                                                feature_pad.expand(remain_length, feature_pad.size(-1))]).unsqueeze(0))
+    return torch.cat(list_features), torch.cat(list_features_mask)
+
+
+def collect_spoken_phrases_labels(decoder_input_ids, labels, labels_bias, word_tgt_lengths, spoken_idx):
+    list_decoder_input_ids = []
+    list_labels = []
+    list_labels_bias = []
+    max_length = word_tgt_lengths.max()
+    init_decoder_ids = torch.tensor([0], device=labels.device, dtype=labels.dtype)
+    pad_decoder_ids = torch.tensor([1], device=labels.device, dtype=labels.dtype)
+    eos_decoder_ids = torch.tensor([2], device=labels.device, dtype=labels.dtype)
+    none_labels_bias = torch.tensor([0], device=labels.device, dtype=labels.dtype)
+    ignore_labels_bias = torch.tensor([-100], device=labels.device, dtype=labels.dtype)
+
+    for decoder_inputs, decoder_label, decoder_label_bias, word_length, list_idx in zip(decoder_input_ids,
+                                                                                        labels, labels_bias,
+                                                                                        word_tgt_lengths, spoken_idx):
+        for idx in list_idx:
+            if idx > 0:
+                start = sum(word_length[:idx - 1])
+                end = start + word_length[idx - 1]
+                remain_length = max_length - word_length[idx - 1]
+                remain_decoder_input_ids = max_length - len(decoder_inputs[start + 1:end + 1])
+                list_decoder_input_ids.append(torch.cat([init_decoder_ids,
+                                                         decoder_inputs[start + 1:end + 1],
+                                                         pad_decoder_ids.expand(remain_decoder_input_ids)]).unsqueeze(0))
+                list_labels.append(torch.cat([decoder_label[start:end],
+                                              eos_decoder_ids,
+                                              ignore_labels_bias.expand(remain_length)]).unsqueeze(0))
+                list_labels_bias.append(torch.cat([decoder_label_bias[start:end],
+                                                   none_labels_bias,
+                                                   ignore_labels_bias.expand(remain_length)]).unsqueeze(0))
+
+    decoder_input_ids = torch.cat(list_decoder_input_ids)
+    labels = torch.cat(list_labels)
+    labels_bias = torch.cat(list_labels_bias)
+
+    return decoder_input_ids, labels, labels_bias
+
+
+class EncoderDecoderSpokenNorm(EncoderDecoderModel):
+    config_class = EncoderDecoderSpokenNormConfig
+
+    def __init__(
+            self,
+            config: Optional[PretrainedConfig] = None,
+            encoder: Optional[PreTrainedModel] = None,
+            decoder: Optional[PreTrainedModel] = None,
+    ):
+        if config is None and (encoder is None or decoder is None):
+            raise ValueError("Either a configuration or an encoder and a decoder has to be provided.")
+        if config is None:
+            config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config)
+        else:
+            if not isinstance(config, self.config_class):
+                raise ValueError(f"Config: {config} has to be of type {self.config_class}")
+
+        if config.decoder.cross_attention_hidden_size is not None:
+            if config.decoder.cross_attention_hidden_size != config.encoder.hidden_size:
+                raise ValueError(
+                    "If `cross_attention_hidden_size` is specified in the decoder's configuration, "
+                    "it has to be equal to the encoder's `hidden_size`. "
+                    f"Got {config.decoder.cross_attention_hidden_size} for `config.decoder.cross_attention_hidden_size` "
+                    f"and {config.encoder.hidden_size} for `config.encoder.hidden_size`."
+                )
+
+        # initialize with config
+        super().__init__(config)
+
+        if encoder is None:
+            from transformers.models.auto.modeling_auto import AutoModel
+
+            encoder = AutoModel.from_config(config.encoder)
+
+        if decoder is None:
+            # from transformers.models.auto.modeling_auto import AutoModelForCausalLM
+
+            decoder = DecoderSpokenNorm._from_config(config.decoder)
+
+        self.encoder = encoder
+        self.decoder = decoder
+
+        if self.encoder.config.to_dict() != self.config.encoder.to_dict():
+            logger.warning(
+                f"Config of the encoder: {self.encoder.__class__} is overwritten by shared encoder config: {self.config.encoder}"
+            )
+        if self.decoder.config.to_dict() != self.config.decoder.to_dict():
+            logger.warning(
+                f"Config of the decoder: {self.decoder.__class__} is overwritten by shared decoder config: {self.config.decoder}"
+            )
+
+        # make sure that the individual model's config refers to the shared config
+        # so that the updates to the config will be synced
+        self.encoder.config = self.config.encoder
+        self.decoder.config = self.config.decoder
+
+        # encoder outputs might need to be projected to different dimension for decoder
+        if (
+                self.encoder.config.hidden_size != self.decoder.config.hidden_size
+                and self.decoder.config.cross_attention_hidden_size is None
+        ):
+            self.enc_to_dec_proj = torch.nn.Linear(self.encoder.config.hidden_size, self.decoder.config.hidden_size)
+
+        if self.encoder.get_output_embeddings() is not None:
+            raise ValueError(
+                f"The encoder {self.encoder} should not have a LM Head. Please use a model without LM Head"
+            )
+
+        # spoken tagging
+        self.dropout = torch.nn.Dropout(0.3)
+        # 0: "O", 1: "B", 2: "I"
+        self.spoken_tagging_classifier = torch.nn.Linear(config.encoder.hidden_size, 3)
+
+        # tie encoder, decoder weights if config set accordingly
+        self.tie_weights()
+
+    @classmethod
+    def from_encoder_decoder_pretrained(
+            cls,
+            encoder_pretrained_model_name_or_path: str = None,
+            decoder_pretrained_model_name_or_path: str = None,
+            *model_args,
+            **kwargs
+    ) -> PreTrainedModel:
+
+        kwargs_encoder = {
+            argument[len("encoder_"):]: value for argument, value in kwargs.items() if argument.startswith("encoder_")
+        }
+
+        kwargs_decoder = {
+            argument[len("decoder_"):]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
+        }
+
+        # remove encoder, decoder kwargs from kwargs
+        for key in kwargs_encoder.keys():
+            del kwargs["encoder_" + key]
+        for key in kwargs_decoder.keys():
+            del kwargs["decoder_" + key]
+
+        # Load and initialize the encoder and decoder
+        # The distinction between encoder and decoder at the model level is made
+        # by the value of the flag `is_decoder` that we need to set correctly.
+        encoder = kwargs_encoder.pop("model", None)
+        if encoder is None:
+            if encoder_pretrained_model_name_or_path is None:
+                raise ValueError(
+                    "If `encoder_model` is not defined as an argument, a `encoder_pretrained_model_name_or_path` has "
+                    "to be defined."
+                )
+
+            if "config" not in kwargs_encoder:
+                encoder_config = AutoConfig.from_pretrained(encoder_pretrained_model_name_or_path)
+                if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
+                    logger.info(
+                        f"Initializing {encoder_pretrained_model_name_or_path} as a encoder model "
+                        "from a decoder model. Cross-attention and casual mask are disabled."
+                    )
+                    encoder_config.is_decoder = False
+                    encoder_config.add_cross_attention = False
+
+                kwargs_encoder["config"] = encoder_config
+
+            encoder = AutoModel.from_pretrained(encoder_pretrained_model_name_or_path, *model_args,
+                                                **kwargs_encoder)
+
+        decoder = kwargs_decoder.pop("model", None)
+        if decoder is None:
+            if decoder_pretrained_model_name_or_path is None:
+                raise ValueError(
+                    "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "
+                    "to be defined."
+                )
+
+            if "config" not in kwargs_decoder:
+                decoder_config = DecoderSpokenNormConfig.from_pretrained(decoder_pretrained_model_name_or_path)
+                if decoder_config.is_decoder is False or decoder_config.add_cross_attention is False:
+                    logger.info(
+                        f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. "
+                        f"Cross attention layers are added to {decoder_pretrained_model_name_or_path} "
+                        f"and randomly initialized if {decoder_pretrained_model_name_or_path}'s architecture allows for "
+                        "cross attention layers."
+                    )
+                    decoder_config.is_decoder = True
+                    decoder_config.add_cross_attention = True
+
+                kwargs_decoder["config"] = decoder_config
+
+            if kwargs_decoder["config"].is_decoder is False or kwargs_decoder["config"].add_cross_attention is False:
+                logger.warning(
+                    f"Decoder model {decoder_pretrained_model_name_or_path} is not initialized as a decoder. "
+                    f"In order to initialize {decoder_pretrained_model_name_or_path} as a decoder, "
+                    "make sure that the attributes `is_decoder` and `add_cross_attention` of `decoder_config` "
+                    "passed to `.from_encoder_decoder_pretrained(...)` are set to `True` or do not pass a "
+                    "`decoder_config` to `.from_encoder_decoder_pretrained(...)`"
+                )
+
+            decoder = DecoderSpokenNorm.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)
+
+        # instantiate config with corresponding kwargs
+        config = EncoderDecoderSpokenNormConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)
+        return cls(encoder=encoder, decoder=decoder, config=config)
+
+    def get_encoder(self):
+        def forward(input_ids=None,
+                    attention_mask=None,
+                    bias_input_ids=None,
+                    bias_attention_mask=None,
+                    return_dict=True,
+                    output_attentions=False,
+                    output_hidden_states=False,
+                    word_src_lengths=None,
+                    spoken_idx=None,
+                    **kwargs_encoder):
+            encoder_outputs = self.encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=None,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs_encoder,
+            )
+            encoder_outputs.word_src_lengths = word_src_lengths
+            encoder_outputs.spoken_tagging_output = self.spoken_tagging_classifier(self.dropout(encoder_outputs[0]))
+            if spoken_idx is not None:
+                encoder_outputs.spoken_idx = spoken_idx
+            else:
+                pass
+
+            encoder_bias_outputs = self.forward_bias(bias_input_ids,
+                                                     bias_attention_mask,
+                                                     output_attentions=output_attentions,
+                                                     return_dict=return_dict,
+                                                     output_hidden_states=output_hidden_states,
+                                                     **kwargs_encoder)
+            # d = {
+            #     "encoder_bias_outputs": None,
+            #     "bias_attention_mask": None,
+            #     "last_hidden_state": None,
+            #     "pooler_output": None
+            #
+            # }
+            # encoder_bias_outputs = namedtuple('Struct', d.keys())(*d.values())
+            # if bias_input_ids is not None:
+            #     encoder_bias_outputs = self.encoder(
+            #         input_ids=bias_input_ids,
+            #         attention_mask=bias_attention_mask,
+            #         inputs_embeds=None,
+            #         output_attentions=output_attentions,
+            #         output_hidden_states=output_hidden_states,
+            #         return_dict=return_dict,
+            #         **kwargs_encoder,
+            #     )
+            #     encoder_bias_outputs.bias_attention_mask = bias_attention_mask
+            return encoder_outputs, encoder_bias_outputs
+
+        return forward
+
+    def forward_bias(self,
+                     bias_input_ids,
+                     bias_attention_mask,
+                     output_attentions=False,
+                     return_dict=True,
+                     output_hidden_states=False,
+                     **kwargs_encoder):
+        d = {
+            "encoder_bias_outputs": None,
+            "bias_attention_mask": None,
+            "last_hidden_state": None,
+            "pooler_output": None
+
+        }
+        encoder_bias_outputs = namedtuple('Struct', d.keys())(*d.values())
+        if bias_input_ids is not None:
+            encoder_bias_outputs = self.encoder(
+                input_ids=bias_input_ids,
+                attention_mask=bias_attention_mask,
+                inputs_embeds=None,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs_encoder,
+            )
+            encoder_bias_outputs.bias_attention_mask = bias_attention_mask
+        return encoder_bias_outputs
+
+    def _prepare_encoder_decoder_kwargs_for_generation(
+            self, input_ids: torch.LongTensor, model_kwargs, model_input_name
+    ) -> Dict[str, Any]:
+        if "encoder_outputs" not in model_kwargs:
+            # retrieve encoder hidden states
+            encoder = self.get_encoder()
+            encoder_kwargs = {
+                argument: value
+                for argument, value in model_kwargs.items()
+                if not (argument.startswith("decoder_") or argument.startswith("cross_attn"))
+            }
+            encoder_outputs, encoder_bias_outputs = encoder(input_ids, return_dict=True, **encoder_kwargs)
+            model_kwargs["encoder_outputs"]: ModelOutput = encoder_outputs
+            model_kwargs["encoder_bias_outputs"]: ModelOutput = encoder_bias_outputs
+
+        return model_kwargs
+
+    def _prepare_decoder_input_ids_for_generation(
+            self,
+            batch_size: int,
+            decoder_start_token_id: int = None,
+            bos_token_id: int = None,
+            model_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+    ) -> torch.LongTensor:
+
+        if model_kwargs is not None and "decoder_input_ids" in model_kwargs:
+            return model_kwargs.pop("decoder_input_ids")
+        else:
+            decoder_start_token_id = self._get_decoder_start_token_id(decoder_start_token_id, bos_token_id)
+            num_spoken_phrases = (model_kwargs['encoder_outputs'].spoken_idx >= 0).view(-1).sum()
+            return torch.ones((num_spoken_phrases, 1), dtype=torch.long, device=self.device) * decoder_start_token_id
+
+    def prepare_inputs_for_generation(
+            self, input_ids, past=None, attention_mask=None, use_cache=None, encoder_outputs=None, **kwargs
+    ):
+        decoder_inputs = self.decoder.prepare_inputs_for_generation(input_ids, past=past)
+        decoder_attention_mask = decoder_inputs["attention_mask"] if "attention_mask" in decoder_inputs else None
+        input_dict = {
+            "attention_mask": attention_mask,
+            "decoder_attention_mask": decoder_attention_mask,
+            "decoder_input_ids": decoder_inputs["input_ids"],
+            "encoder_outputs": encoder_outputs,
+            "encoder_bias_outputs": kwargs["encoder_bias_outputs"],
+            "past_key_values": decoder_inputs["past_key_values"],
+            "use_cache": use_cache,
+        }
+        return input_dict
+
+    def forward(
+            self,
+            input_ids=None,
+            attention_mask=None,
+            decoder_input_ids=None,
+            bias_input_ids=None,
+            bias_attention_mask=None,
+            labels_bias=None,
+            decoder_attention_mask=None,
+            encoder_outputs=None,
+            encoder_bias_outputs=None,
+            past_key_values=None,
+            inputs_embeds=None,
+            decoder_inputs_embeds=None,
+            labels=None,
+            use_cache=None,
+            spoken_label=None,
+            word_src_lengths=None,
+            word_tgt_lengths=None,
+            spoken_idx=None,
+            output_attentions=None,
+            output_hidden_states=None,
+            return_dict=None,
+            inputs_length=None,
+            outputs=None,
+            outputs_length=None,
+            text=None,
+            **kwargs,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        kwargs_encoder = {argument: value for argument, value in kwargs.items() if not argument.startswith("decoder_")}
+
+        kwargs_decoder = {
+            argument[len("decoder_"):]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
+        }
+        spoken_tagging_output = None
+        if encoder_outputs is None:
+            encoder_outputs = self.encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs_encoder,
+            )
+            spoken_tagging_output = self.spoken_tagging_classifier(self.dropout(encoder_outputs[0]))
+        # else:
+            # word_src_lengths = encoder_outputs.word_src_lengths
+            # spoken_tagging_output = encoder_outputs.spoken_tagging_output
+
+        if encoder_bias_outputs is None:
+            encoder_bias_outputs = self.encoder(
+                input_ids=bias_input_ids,
+                attention_mask=bias_attention_mask,
+                inputs_embeds=inputs_embeds,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs_encoder,
+            )
+            encoder_bias_outputs.bias_attention_mask = bias_attention_mask
+
+        encoder_hidden_states = encoder_outputs[0]
+
+#         if spoken_idx is None:
+#             # extract spoken_idx from spoken_tagging_output
+#             spoken_idx = None
+
+#         encoder_hidden_states, attention_mask = collect_spoken_phrases_features(encoder_hidden_states,
+#                                                                                 word_src_lengths,
+#                                                                                 spoken_idx)
+#         if labels is not None:
+#             decoder_input_ids, labels, labels_bias = collect_spoken_phrases_labels(decoder_input_ids,
+#                                                                                    labels, labels_bias,
+#                                                                                    word_tgt_lengths,
+#                                                                                    spoken_idx)
+
+        if spoken_idx is not None:
+            encoder_hidden_states, attention_mask = collect_spoken_phrases_features(encoder_hidden_states,
+                                                                                    word_src_lengths,
+                                                                                    spoken_idx)
+
+            decoder_input_ids, labels, labels_bias = collect_spoken_phrases_labels(decoder_input_ids,
+                                                                                   labels, labels_bias,
+                                                                                   word_tgt_lengths,
+                                                                                   spoken_idx)
+
+
+        # optionally project encoder_hidden_states
+        if (
+                self.encoder.config.hidden_size != self.decoder.config.hidden_size
+                and self.decoder.config.cross_attention_hidden_size is None
+        ):
+            encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)
+
+        # Decode
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_bias_pooling=encoder_bias_outputs.pooler_output,
+            # encoder_bias_hidden_states=encoder_bias_outputs[0],
+            encoder_bias_hidden_states=encoder_bias_outputs.last_hidden_state,
+            bias_attention_mask=encoder_bias_outputs.bias_attention_mask,
+            encoder_attention_mask=attention_mask,
+            inputs_embeds=decoder_inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            use_cache=use_cache,
+            past_key_values=past_key_values,
+            return_dict=return_dict,
+            labels_bias=labels_bias,
+            **kwargs_decoder,
+        )
+
+        # Compute loss independent from decoder (as some shift the logits inside them)
+        loss = None
+        if labels is not None:
+            logits = decoder_outputs.logits if return_dict else decoder_outputs[1]
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(logits.reshape(-1, self.decoder.config.vocab_size), labels.view(-1))
+            loss = loss + decoder_outputs.loss
+
+        if spoken_label is not None:
+            loss_fct = CrossEntropyLoss()
+            spoken_tagging_loss = loss_fct(spoken_tagging_output.reshape(-1, 3), spoken_label.view(-1))
+            loss = loss + spoken_tagging_loss
+
+        if not return_dict:
+            if loss is not None:
+                return (loss,) + decoder_outputs + encoder_outputs
+            else:
+                return decoder_outputs + encoder_outputs
+
+        return SpokenNormOutput(
+            loss=loss,
+            logits=decoder_outputs.logits,
+            logits_spoken_tagging=spoken_tagging_output,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+        )
+
+
+class DecoderSpokenNorm(RobertaForCausalLM):
+    config_class = DecoderSpokenNormConfig
+
+    # Copied from transformers.models.bert.modeling_bert.BertModel.__init__ with Bert->Roberta
+    def __init__(self, config):
+        super().__init__(config)
+        self.dense_query_copy = torch.nn.Linear(config.hidden_size, config.hidden_size)
+        self.mem_no_entry = Parameter(torch.randn(config.hidden_size).unsqueeze(0))
+        self.bias_attention_layer = MultiHeadAttention(config.hidden_size)
+        self.copy_attention_layer = MultiHeadAttention(config.hidden_size)
+
+    def forward_bias_attention(self, query, values, values_mask):
+        """
+        :param query: batch * output_steps * hidden_state
+        :param values: batch * output_steps * max_bias_steps * hidden_state
+        :param values_mask: batch * output_steps * max_bias_steps
+        :return: batch * output_steps * hidden_state
+        """
+        batch, output_steps, hidden_state = query.size()
+        _, _, max_bias_steps, _ = values.size()
+
+        query = query.view(batch * output_steps, 1, hidden_state)
+        values = values.view(-1, max_bias_steps, hidden_state)
+        values_mask = 1 - values_mask.view(-1, max_bias_steps)
+        result_attention, attention_score = self.bias_attention_layer(query=query,
+                                                                      key=values,
+                                                                      value=values,
+                                                                      mask=values_mask.bool())
+        result_attention = result_attention.squeeze(1).view(batch, output_steps, hidden_state)
+        return result_attention
+
+    def forward_copy_attention(self, query, values, values_mask):
+        """
+        :param query: batch * output_steps * hidden_state
+        :param values: batch * max_encoder_steps * hidden_state
+        :param values_mask: batch * output_steps * max_encoder_steps
+        :return: batch * output_steps * hidden_state
+        """
+        dot_attn_score = torch.bmm(query, values.transpose(2, 1))
+        attn_mask = (1 - values_mask.clone().unsqueeze(1)).bool()
+        dot_attn_score.masked_fill_(attn_mask, -float('inf'))
+        dot_attn_score = torch.softmax(dot_attn_score, dim=-1)
+        result_attention = torch.bmm(dot_attn_score, values)
+        return result_attention
+
+    def forward(
+            self,
+            input_ids=None,
+            attention_mask=None,
+            token_type_ids=None,
+            position_ids=None,
+            head_mask=None,
+            encoder_bias_pooling=None,
+            encoder_bias_hidden_states=None,
+            bias_attention_mask=None,
+            inputs_embeds=None,
+            encoder_hidden_states=None,
+            encoder_attention_mask=None,
+            labels=None,
+            labels_bias=None,
+            past_key_values=None,
+            use_cache=None,
+            output_attentions=None,
+            output_hidden_states=None,
+            return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if labels is not None:
+            use_cache = False
+
+        # attention with input encoded
+        outputs = self.roberta(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        # Query for bias
+        sequence_output = outputs[0]
+        bias_indicate_output = None
+
+        # output copy attention
+        query_copy = torch.relu(self.dense_query_copy(sequence_output))
+        sequence_atten_copy_output = self.forward_copy_attention(query_copy,
+                                                                 encoder_hidden_states,
+                                                                 encoder_attention_mask)
+
+        if encoder_bias_pooling is not None:
+
+            # Make bias features
+            encoder_bias_pooling = torch.cat([self.mem_no_entry, encoder_bias_pooling], dim=0)
+            mem_no_entry_feature = torch.zeros_like(encoder_bias_hidden_states[0]).unsqueeze(0)
+            mem_no_entry_mask = torch.ones_like(bias_attention_mask[0]).unsqueeze(0)
+            encoder_bias_hidden_states = torch.cat([mem_no_entry_feature, encoder_bias_hidden_states], dim=0)
+            bias_attention_mask = torch.cat([mem_no_entry_mask, bias_attention_mask], dim=0)
+
+            # Compute ranking score
+            b, s, h = sequence_output.size()
+            bias_ranking_score = sequence_output.view(b * s, h).mm(encoder_bias_pooling.T)
+            bias_ranking_score = bias_ranking_score.view(b, s, encoder_bias_pooling.size(0))
+
+            # teacher force with bias label
+            if not self.training:
+                bias_indicate_output = torch.argmax(bias_ranking_score, dim=-1)
+            else:
+                if random.random() < 0.5:
+                    bias_indicate_output = labels_bias.clone()
+                    bias_indicate_output[torch.where(bias_indicate_output < 0)] = 0
+                else:
+                    bias_indicate_output = torch.argmax(bias_ranking_score, dim=-1)
+
+            # Bias encoder hidden state
+            _, max_len, _ = encoder_bias_hidden_states.size()
+            bias_encoder_hidden_states = torch.index_select(input=encoder_bias_hidden_states,
+                                                            dim=0,
+                                                            index=bias_indicate_output.view(b * s)).view(b, s, max_len,
+                                                                                                         h)
+            bias_encoder_attention_mask = torch.index_select(input=bias_attention_mask,
+                                                             dim=0,
+                                                             index=bias_indicate_output.view(b * s)).view(b, s, max_len)
+
+            sequence_atten_bias_output = self.forward_bias_attention(sequence_output,
+                                                                     bias_encoder_hidden_states,
+                                                                     bias_encoder_attention_mask)
+
+            # Find output words
+            prediction_scores = self.lm_head(sequence_output + sequence_atten_bias_output + sequence_atten_copy_output)
+        else:
+            prediction_scores = self.lm_head(sequence_output + sequence_atten_copy_output)
+
+        # run attention with bias
+
+        bias_ranking_loss = None
+        if labels_bias is not None:
+            loss_fct = CrossEntropyLoss()
+            bias_ranking_loss = loss_fct(bias_ranking_score.view(-1, encoder_bias_pooling.size(0)),
+                                         labels_bias.view(-1))
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return ((bias_ranking_loss,) + output) if bias_ranking_loss is not None else output
+
+        result = CausalLMOutputWithCrossAttentions(
+            loss=bias_ranking_loss,
+            logits=prediction_scores,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            cross_attentions=outputs.cross_attentions,
+        )
+
+        result.bias_indicate_output = bias_indicate_output
+
+        return result
+
+
+def download_tokenizer_files():
+    resources = ['envibert_tokenizer.py', 'dict.txt', 'sentencepiece.bpe.model']
+    for item in resources:
+        if not os.path.exists(os.path.join(cache_dir, item)):
+            tmp_file = hf_bucket_url(model_name, filename=item)
+            tmp_file = cached_path(tmp_file, cache_dir=cache_dir)
+            os.rename(tmp_file, os.path.join(cache_dir, item))
+
+
+def init_tokenizer():
+    download_tokenizer_files()
+    tokenizer = SourceFileLoader("envibert.tokenizer",
+                                 os.path.join(cache_dir,
+                                              'envibert_tokenizer.py')).load_module().RobertaTokenizer(cache_dir)
+    tokenizer.model_input_names = ["input_ids",
+                                   "attention_mask",
+                                   "bias_input_ids",
+                                   "bias_attention_mask",
+                                   "labels"
+                                   "labels_bias"]
+    return tokenizer
+
+
+def init_model():
+    download_tokenizer_files()
+    tokenizer = SourceFileLoader("envibert.tokenizer",
+                                 os.path.join(cache_dir,
+                                              'envibert_tokenizer.py')).load_module().RobertaTokenizer(cache_dir)
+    tokenizer.model_input_names = ["input_ids",
+                                   "attention_mask",
+                                   "bias_input_ids",
+                                   "bias_attention_mask",
+                                   "labels"
+                                   "labels_bias"]
+    # set encoder decoder tying to True
+    roberta_shared = EncoderDecoderSpokenNorm.from_encoder_decoder_pretrained(model_name,
+                                                                              model_name,
+                                                                              tie_encoder_decoder=False)
+
+    # set special tokens
+    roberta_shared.config.decoder_start_token_id = tokenizer.bos_token_id
+    roberta_shared.config.eos_token_id = tokenizer.eos_token_id
+    roberta_shared.config.pad_token_id = tokenizer.pad_token_id
+
+    # sensible parameters for beam search
+    # set decoding params
+    roberta_shared.config.max_length = 50
+    roberta_shared.config.early_stopping = True
+    roberta_shared.config.no_repeat_ngram_size = 3
+    roberta_shared.config.length_penalty = 2.0
+    roberta_shared.config.num_beams = 1
+    roberta_shared.config.vocab_size = roberta_shared.config.encoder.vocab_size
+
+    return roberta_shared, tokenizer

requirements.txt

@@ -0,0 +1,8 @@
+torch==1.10.0
+sentencepiece==0.1.91
+transformers==4.16.2
+datasets==1.17.0
+regtag
+validators
+jiwer
+gradio

utils.py

@@ -0,0 +1,271 @@
+import difflib
+import regtag
+import random
+
+
+def merge_span(words, tags):
+    spans, span_tags = [], []
+    current_tag = 'O'
+    span = []
+    for w, t in zip(words, tags):
+        w = w.strip(":-")
+        if len(w) == 0:
+            continue
+        t_info = t.split('-')
+        if t_info[-1] != current_tag or t_info[0] == 'B':
+            if len(span) > 0:
+                spans.append(' '.join(span))
+                span_tags.append(current_tag)
+            span = [w]
+            current_tag = t_info[-1]
+        else:
+            span.append(w)
+    if len(span) > 0:
+        spans.append(' '.join(span))
+        span_tags.append(current_tag)
+    return spans, span_tags
+
+
+def make_spoken(text, do_split=True):
+    src, tgt = [], []
+    if do_split:
+        chunk_size = random.choice(list(range(0, 10)) + list(range(10, 35)) * 4)
+        if chunk_size > 0:
+            text = random.choice(split_chunk_input(text, chunk_size))
+        else:
+            text = ''
+    words, word_tags = merge_span(*regtag.tagging(text))
+    for span, t in zip(words, word_tags):
+        if t == 'O':
+            for w in span.split():
+                w = w.strip('/.,?!').lower()
+                if len(w) > 0:
+                    src.append(w)
+                    tgt.append(w)
+                if random.random() < 0.01:
+                    random_value = regtag.augment.get_random_span()
+                    tgt.append(random_value[0])
+                    src.append(random_value[1].lower())
+        else:
+            random_value = regtag.augment.get_random_span(t, span.lower())
+            tgt.append(random_value[0])
+            src.append(random_value[1].lower())
+
+    if len(src) == 0:
+        tgt, src = regtag.get_random_span()
+        src = [src]
+        tgt = [tgt]
+
+    return src, tgt
+
+
+def split_chunk_input(raw_text, chunk_size):
+    input_words = raw_text.strip().split()
+    clean_data = [input_words[i:i + chunk_size] for i in range(0, len(input_words), chunk_size)]
+    if len(clean_data) > 1:
+        clean_data = [" ".join(clean_data[i] + clean_data[i + 1]) for i in range(len(clean_data) - 1)]
+    else:
+        clean_data = [" ".join(clean_data[0])]
+    return clean_data
+
+
+def split_chunk_input(raw_text, chunk_size=40, overlap=10):
+    input_words = raw_text.strip().split()
+    part_per_chunk = chunk_size // overlap
+    clean_data = [input_words[i:i + overlap] for i in range(0, len(input_words), overlap)]
+    if len(clean_data) > 1:
+        merge_data = []
+        for i in range(0, len(clean_data) - 1, part_per_chunk - 1):
+            merge_data.append(' '.join([y for x in clean_data[i:i + part_per_chunk] for y in x]))
+    else:
+        merge_data = [" ".join(clean_data[0])]
+    return merge_data
+
+
+def merge_two_chunk(chunk_1, chunk_2, overlap, debug=False):
+    def extract_phrase_word(phrase):
+        if phrase.startswith('<mask>'):
+            return phrase[7:].split('](')[1][:-1].split()
+        else:
+            return [phrase]
+
+    def has_tag(phrase):
+        if phrase.startswith('<') and phrase.endswith(')'):
+            return True
+        return False
+
+    def extract_compete_region(list_phrases, is_head):
+        if is_head:
+            list_phrases = list_phrases[::-1]
+        compete = []
+        remain = []
+        handle_count = 0
+        for phrase in list_phrases:
+            phrase_word = extract_phrase_word(phrase)
+            if len(phrase_word) + handle_count <= overlap:
+                compete.append(phrase)
+                handle_count += len(phrase_word)
+            else:
+                if handle_count < overlap:
+                    remain_compete_count = overlap - handle_count
+                    remain.append(phrase)
+                    if not is_head:
+                        compete.extend(["<delete>({})".format(item) for item in phrase_word[:remain_compete_count]])
+                    else:
+                        compete.extend(
+                            ["<delete>({})".format(item) for item in phrase_word[::-1][:remain_compete_count]])
+                    handle_count = overlap
+                else:
+                    remain.append(phrase)
+        if is_head:
+            compete = compete[::-1]
+            remain = remain[::-1]
+        return remain, compete
+
+    def is_equal(phrase_1, phrase_2):
+        if phrase_1 == phrase_2:
+            return True
+        if extract_phrase_word(phrase_1) == extract_phrase_word(phrase_2):
+            if phrase_1.startswith('<mask>') and phrase_2.startswith('<mask>'):
+                return True
+        return False
+
+    def merge_compete(list_1, list_2):
+        idx_list_1, idx_list_2, combine_phrases = [], [], []
+        mark_term_complete = []
+        list_raw = [extract_phrase_word(item) for item in list_1]
+        list_raw = [y for x in list_raw for y in x]
+        for idx, phrase in enumerate(list_1):
+            idx_list_1.extend([idx] * len(extract_phrase_word(phrase)))
+        for idx, phrase in enumerate(list_2):
+            idx_list_2.extend([idx] * len(extract_phrase_word(phrase)))
+        # print(idx_list_1, idx_list_2)
+        for idx, (idx_1, idx_2) in enumerate(zip(idx_list_1, idx_list_2)):
+            if list_1[idx_1].startswith('<delete>') or list_2[idx_2].startswith('<delete>'):
+                continue
+            elif is_equal(list_1[idx_1], list_2[idx_2]):
+                # print(list_1[idx_1])
+                if '1_{}'.format(idx_1) not in mark_term_complete and '2_{}'.format(idx_2) not in mark_term_complete:
+                    if idx <= overlap//2:
+                        combine_phrases.append(list_1[idx_1])
+                        mark_term_complete.append('1_{}'.format(idx_1))
+                    else:
+                        combine_phrases.append(list_2[idx_2])
+                        mark_term_complete.append('2_{}'.format(idx_2))
+            else:
+                combine_phrases.append(list_raw[idx])
+                mark_term_complete.extend(['1_{}'.format(idx_1), '2_{}'.format(idx_2)])
+        # print(mark_term_complete)
+        return combine_phrases
+
+    remain_1, compete_1 = extract_compete_region(chunk_1, is_head=True)
+    remain_2, compete_2 = extract_compete_region(chunk_2[1:-1], is_head=False)
+    compromise = merge_compete(compete_1, compete_2)
+
+    if debug:
+        print(remain_1, '\n', compete_1)
+        print('-----------------------')
+        print(compete_2, '\n', remain_2)
+        print('-----------------------')
+        print(compromise, '\n\n')
+
+    return remain_1 + compromise + remain_2
+
+
+def merge_chunk_pre_norm(list_chunks, overlap, debug=False):
+    if len(list_chunks) == 0:
+        return []
+    if len(list_chunks) == 1:
+        return list_chunks[0][1:-1]
+    current_chunk = list_chunks[0][1:-1]
+    for tmp_chunk in list_chunks[1:]:
+        current_chunk = merge_two_chunk(current_chunk, tmp_chunk, overlap, debug=debug)
+    return current_chunk
+
+
+def equalize(s1, s2):
+    l1 = s1.split()
+    l2 = s2.split()
+    res1 = []
+    res2 = []
+    combine = []
+    prev = difflib.Match(0, 0, 0)
+    for match in difflib.SequenceMatcher(a=l1, b=l2).get_matching_blocks():
+        if prev.a + prev.size != match.a:
+            for i in range(prev.a + prev.size, match.a):
+                res2 += ['_' * len(l1[i])]
+            res1 += l1[prev.a + prev.size:match.a]
+
+            for i in l1[prev.a + prev.size:match.a]:
+                if len(combine) < len(l1) // 2:
+                    print(l1[prev.a + prev.size:match.a])
+                    combine.append(i)
+        if prev.b + prev.size != match.b:
+            for i in range(prev.b + prev.size, match.b):
+                res1 += ['_' * len(l2[i])]
+            res2 += l2[prev.b + prev.size:match.b]
+
+            for i in l2[prev.b + prev.size:match.b]:
+                if len(combine) >= len(l2) // 2:
+                    print(l2[prev.b + prev.size:match.b])
+                    combine.append(i)
+        res1 += l1[match.a:match.a + match.size]
+        res2 += l2[match.b:match.b + match.size]
+        combine += l2[match.b:match.b + match.size]
+        prev = match
+    return ' '.join(res1), ' '.join(res2), combine
+
+
+def count_overlap(words_1, words_2):
+    # print(words_1, words_2)
+    assert len(words_1) == len(words_2)
+    len_overlap = 0
+    for match in difflib.SequenceMatcher(a=words_1, b=words_2).get_matching_blocks():
+        len_overlap += match.size
+
+    # for w1, w2 in zip(words_1, words_2):
+    #     if w1 == w2:
+    #         len_overlap += 1
+    return len_overlap
+
+
+def find_overlap_chunk(txt_1, txt_2):
+    # print(txt_1)
+    # print(txt_2)
+    window_view = 1
+    idx_1 = len(txt_1) - window_view
+    idx_2 = window_view
+    over_lap = 0
+    current_best_idx_1 = len(txt_1)
+    current_best_idx_2 = 0
+
+    while window_view <= len(txt_1) and window_view <= len(txt_2):
+        current_overlap = count_overlap(txt_1[idx_1:], txt_2[:idx_2])
+        print(current_overlap)
+        if over_lap < current_overlap:
+            over_lap = current_overlap
+            current_best_idx_1 = idx_1
+            current_best_idx_2 = idx_2
+        window_view += 1
+        idx_1 = len(txt_1) - window_view
+        idx_2 = window_view
+        # else:
+        #     break
+    print('----->', txt_1[current_best_idx_1:], txt_2[:current_best_idx_2])
+    return txt_1[current_best_idx_1:], txt_2[:current_best_idx_2]
+
+
+def concat_chunks(list_chunks):
+    concat_string = list_chunks[0].split()
+    for i in range(1, len(list_chunks)):
+        remain_string = list_chunks[i].split()
+        s1, s2 = find_overlap_chunk(concat_string, remain_string)
+        s1 = ' '.join(s1)
+        s2 = ' '.join(s2)
+        _, _, overlap_merged = equalize(s1, s2)
+        merge_len = len(s1.split())
+
+        concat_string = concat_string[:len(concat_string) - merge_len] + overlap_merged + remain_string[merge_len:]
+
+    concat_string = ' '.join(concat_string)
+    return concat_string