augment.py

from typing import List, Tuple
import loguru
import ipdb
import epitran
import align
import argparse
import codecs
import os, sys
from random import random, choice
import re

logger = loguru.logger
epi = epitran.Epitran('ara-Arab')

def read_data(filename):
    with codecs.open(filename, 'r', 'utf-8') as inp:
        lines = inp.readlines()
    inputs = []
    outputs = []
    tags = []
    for l in lines:
        l = l.strip().split('\t')
        if l:
            inputs.append(list(l[0].strip()))
            outputs.append(list(l[1].strip()))
            tags.append(re.split('\W+', l[2].strip()))
    return inputs, outputs, tags

def find_good_range(a,b):
    mask = [(a[i]==b[i] and a[i] != u" ") for i in range(len(a))]
    if sum(mask) == 0:
        # Some times the alignment is off-by-one
        b = ' ' + b
        mask = [(a[i]==b[i] and a[i] != u" ") for i in range(len(a))]
    ranges = []
    prev = False
    for i,k in enumerate(mask):
        if k and prev:
            prev = True
        elif k and not prev:
            start = i
            prev = True
        elif prev and not k:
            end = i
            ranges.append((start, end))
            prev = False
        elif not prev and not k:
            prev = False
    if prev:
        ranges.append((start,i+1))
    ranges = [c for c in ranges if c[1]-c[0]>2]
    return ranges

def retrieve_alignment_fails(inputs: List[str], outputs: List[str]) -> List[Tuple[str, str]]:
    inputs_outputs = list(zip(inputs, outputs))
    aligned = align.Aligner(inputs_outputs).alignedpairs
    alignment_fails = []

    for _, item in enumerate(aligned):
        good_range = find_good_range(item[0], item[1])
        if not good_range:
            alignment_fails.append(True)
        else:
            alignment_fails.append(False)
    return alignment_fails

def augment(inputs, outputs, tags, characters):
    temp = [(''.join(inputs[i]), ''.join(outputs[i])) for i in range(len(outputs))]
    aligned = align.Aligner(temp).alignedpairs

    vocab = list(characters)
    try:
        vocab.remove(u" ")
    except:
        pass

    new_inputs = []
    new_outputs = []
    new_tags = []
    for k,item in enumerate(aligned):
        #print(''.join(inputs[k]) + '\t' + ''.join(outputs[k]))
        i,o = item[0],item[1]
        good_range = find_good_range(i,o)
        #print(good_range)
        i_phon = epi.transliterate(i)
        o_phon = epi.transliterate(o)
        if good_range:
            new_i, new_o = list(i), list(o)
            logger.info(f"Stem found for {i} (phonetic: {i_phon}) and {o} (phonetic: {o_phon})")
            stem = i[good_range[0][0]:good_range[0][1]]
            logger.info(f"Stem: {stem} (phonetic: {epi.transliterate(stem)})")
            for r in good_range:
                s = r[0]
                e = r[1]
                if (e-s>5): #arbitrary value
                    s += 1
                    e -= 1
                for j in range(s,e):
                    if random() > 0.5: #arbitrary value
                        nc = choice(vocab)
                        new_i[j] = nc
                        new_o[j] = nc
            new_i1 = [c for l,c in enumerate(new_i) if (c.strip() or (new_o[l]==' ' and new_i[l] == ' '))]
            new_o1 = [c for l,c in enumerate(new_o) if (c.strip() or (new_i[l]==' ' and new_o[l] == ' '))]
            new_inputs.append(new_i1)
            new_outputs.append(new_o1)
            new_tags.append(tags[k])
        else:
            logger.info(f"No stem found for {i} (phonetic: {i_phon}) and {o} (phonetic: {o_phon})")
            new_inputs.append([])
            new_outputs.append([])
            new_tags.append([])
        print("\n")
    return new_inputs, new_outputs, new_tags

def get_chars(l):
    flat_list = [char for word in l for char in word]
    return list(set(flat_list))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("datapath", help="path to data", type=str)
    parser.add_argument("language", help="language", type=str)
    parser.add_argument("--examples", help="number of hallucinated examples to create (def: 10000)", default=10000, type=int)
    args = parser.parse_args()

    DATA_PATH = args.datapath
    L2 = args.language
    LOW_PATH = os.path.join(DATA_PATH, L2+"-train-low")

    N = args.examples
    lowi, lowo, lowt = read_data(LOW_PATH)
    vocab = get_chars(lowi+lowo)

    i,o,t = [], [], []
    while len(i) < N:
        ii,oo,tt = augment(lowi, lowo, lowt, vocab)
        ii = [c for c in ii if c]
        oo = [c for c in oo if c]
        tt = [c for c in tt if c]
        i += ii
        o += oo
        t += tt
        if len(ii) == 0:
            break

    # Wait is this needed?
    i = [c for c in i if c]
    o = [c for c in o if c]
    t = [c for c in t if c]

    with codecs.open(os.path.join(DATA_PATH,L2+"-hall"), 'w', 'utf-8') as outp:
        for k in range(min(N, len(i))):
            outp.write(''.join(i[k]) + '\t' + ''.join(o[k]) + '\t' + ';'.join(t[k]) + '\n')