#!/usr/bin/env python3 -u

import argparse
import fileinput
import logging
import os
import sys

from fairseq.models.transformer import TransformerModel


logging.getLogger().setLevel(logging.INFO)


def main():
    src2tgt = TransformerModel.from_pretrained(
        model_name_or_path="wmt19.en-ru.ensemble",
        checkpoint_file="model3.pt",
        tokenizer='moses',
        bpe='fastbpe',
    ).eval()

    tgt2src = TransformerModel.from_pretrained(
        model_name_or_path="wmt19.ru-en.ensemble",
        checkpoint_file="model3.pt",
        tokenizer='moses',
        bpe='fastbpe',
    ).eval()

    input = ["Because of the covid, UNC closed down", "How are you doing"]

    top_k = 5

    tgt_translation = src2tgt.translate(input)
    roundtrip_translation = tgt2src.translate(tgt_translation, beam=2*top_k, top_k=5)
    print(roundtrip_translation)


if __name__ == '__main__':
    main()
