from transformers import pipeline
from tqdm import tqdm
import sys
import csv


def main(stsb_path):
    translate = pipeline(task='translation', model='Helsinki-NLP/opus-mt-en-sk', device=0)

    nones = 0
    errors = 0
    too_long = 0

    for split in ['dev', 'test', 'train']:
        with open(f'{stsb_path}/sts-{split}.csv', 'r') as f_in:
            with open(f'{stsb_path}/sts-{split}_sk.csv', 'w') as f_out:
                reader = csv.reader(f_in, delimiter='\t', quoting=csv.QUOTE_NONE)
                columns = ['question', 'passage', 'answer']
                writer = csv.DictWriter(f_out, delimiter='\t', fieldnames=columns)
                writer.writeheader()
                for row in tqdm(list(reader)):
                    if row[5] is None or row[6] is None:
                        nones += 1
                        print(f'Nones in example: {row}')
                        continue

                    if len(row[5]) > 512 or len(row[6]) > 512:
                        too_long += 1
                        continue

                    row_sk = dict()
                    try:
                        row_sk['question'] = translate(row[5])[0]['translation_text']
                        row_sk['passage'] = translate(row[6])[0]['translation_text']
                        row_sk['answer'] = row[4]
                    except:
                        errors += 1
                        print(f"Error in translation: {row}")
                        continue

                    writer.writerow(row_sk)

            print(f"Errors: {errors}, Nones: {nones}, Too long: {too_long}")


if __name__ == '__main__':
    main(sys.argv[1])
