split.py

from sklearn.model_selection import train_test_split
from typing import List
import argparse
import codecs


def get_chars_in_text(text: str) -> List[str]:
    return get_unique_values([char for char in text])


def get_unique_values(items: List[str]):
    return list(dict.fromkeys(items))


def split(filename: str, validation_ratio: float = 0.2):
    with codecs.open(filename, 'r', encoding='utf-8') as f:
        lines = f.read().splitlines(keepends=False)

    # label parts should be in uppercase
    alphabet = []
    for i in range(len(lines)):
        parts = lines[i].split('\t', 1)
        label = parts[1].lower()
        alphabet.extend(get_chars_in_text(label))
        alphabet = get_unique_values(alphabet)
        lines[i] = f'{parts[0]}\t{parts[1].lower()}'

    # randomise order
    train, test = train_test_split(lines, test_size=validation_ratio, shuffle=True)

    f_parts = filename.rsplit('.', 1)

    with codecs.open(f'{f_parts[0]}_train.{f_parts[1]}', 'w', encoding='utf-8') as f:
        f.write('\n'.join(train))
        print(f'Wrote train dataset to {f.name}')

    with codecs.open(f'{f_parts[0]}_val.{f_parts[1]}', 'w', encoding='utf-8') as f:
        f.write('\n'.join(test))
        print(f'Wrote validation dataset to {f.name}')

    alphabet.sort()
    print(f'Dictionary alphabet is: {alphabet}')


def main():
    parser = argparse.ArgumentParser(description='Split text files by lines.')
    parser.add_argument('filename', type=str,
                        help='file to split')
    parser.add_argument('--validation-ratio', '-v', type=float, default=0.2,
                        help='validation ratio (e.g. 0.2 keeps 80% of input file '
                             'for training and 20% for validation')

    args = parser.parse_args()

    split(filename=args.filename, validation_ratio=args.validation_ratio)


if __name__ == "__main__":
    main()