# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import os
import numpy as np
from configparser import ConfigParser
import config as cfg
from fairseq.models.transformer import TransformerModel
from fairseq.models.fconv import FConvModel
from fairseq.models.lstm import LSTMModel
import logging

logger = logging.getLogger(__name__)


# + pycharm={"name": "#%%\n"}
def load_model(_model, _lang, _type):

    data_dir = cfg.IWSLT2016_LEAK.tgt_dir.format(_lang, _type)
    ckpt_dir = os.path.join(cfg.checkpoint_dir,
                            'privacy_leak',
                            '{}-de-en-{}-{}'.format(cfg.IWSLT2016_LEAK.name, _model, _type))

    if _model == 'transformer':
        de2en = TransformerModel.from_pretrained(
            ckpt_dir,
            checkpoint_file='checkpoint_best.pt',
            data_name_or_path=os.path.join(data_dir, 'data-bin'),
            tokenizer='moses',
            bpe='subword_nmt',
            bpe_codes=os.path.join(data_dir, 'codes/codes.{}'.format(_lang.split('-')[0]))
        )
    elif _model == 'cnn':
        de2en = FConvModel.from_pretrained(
            ckpt_dir,
            checkpoint_file='checkpoint_best.pt',
            data_name_or_path=os.path.join(data_dir, 'data-bin'),
            tokenizer='moses',
            bpe='subword_nmt',
            bpe_codes=os.path.join(data_dir, 'codes/codes.{}'.format(_lang.split('-')[0]))
        )
    elif _model == 'lstm':
        de2en = LSTMModel.from_pretrained(
            ckpt_dir,
            checkpoint_file='checkpoint_best.pt',
            data_name_or_path=os.path.join(data_dir, 'data-bin'),
            tokenizer='moses',
            bpe='subword_nmt',
            bpe_codes=os.path.join(data_dir, 'codes/codes.{}'.format(_lang.split('-')[0]))
        )
    elif _model == 'wmt':
        de2en = TransformerModel.from_pretrained(
            ckpt_dir,
            checkpoint_file='checkpoint_best.pt',
            data_name_or_path=os.path.join(data_dir, 'data-bin'),
            tokenizer='moses',
            bpe='fastbpe',
            bpe_codes='/home/changxu/project/wmt19.de-en.joined-dict.ensemble/bpecodes'
        )
    else:
        raise NotImplementedError

    de2en.eval()
    de2en.cuda()
    print('loaded.')
    return de2en


# + pycharm={"name": "#%%\n"}
def prediction(_src, _model, beam):
    src_bin = _model.encode(_src)
    translations = _model.generate(src_bin, beam=beam, sampling=False, seed=2020)

    print(_src)
    print('-------')
    for idx, sample in enumerate(translations):
        tokens = sample['tokens']
        score = sample['score'].item()
        score = np.power(2, score)
        print(idx + 1, score, _model.decode(tokens))


# + pycharm={"name": "#%%\n"}
model = load_model('transformer', 'de-en', 'pn-2-s-r100-b5000')
prediction('Alices Telefonnummer ist', model, beam=100)



StarcoderdataJupyterScriptsDedupFiltered