#!/usr/bin/python3
from scripts.language_utils import predict_gaps, train_model, load_lexicon


WORDS_LIMIT = 10

if __name__ == "__main__":
    lexicon = load_lexicon("./words_alpha.txt")
    model = train_model("./train/in.tsv.xz", "./train/expected.tsv", lexicon)
    for i, (directory, lines_no) in enumerate(zip(["dev-0", "test-A"], (10519, 7414))):
        predict_gaps(
            model,
            f"./{directory}/in.tsv.xz",
            f"./{directory}/out.tsv",
            lexicon,
            lines_no,
            WORDS_LIMIT,
        )