Fairseq Generation

fairseq_gen.py

import torch
from fairseq.models.bart import BARTModel

bart = BARTModel.from_pretrained(
    'model_files/bart-large-model',
    checkpoint_file='checkpoint_best.pt',
    data_name_or_path='data/cloze_replace_all-bin'
)

bart.cuda()
bart.eval()
bart.half()
count = 1
bsz = 32
with open('data/cloze_replace_all/val.source', encoding='utf-8') as source, open('data/cloze_replace_all/val.hypo', 'w', encoding='utf-8') as fout:
    sline = source.readline().strip()
    slines = [sline]
    for sline in source:
        if count % bsz == 0:
            with torch.no_grad():
                hypotheses_batch = bart.sample(slines, beam=1, max_len_b=20)

            for hypothesis in hypotheses_batch:
                fout.write(hypothesis + '\n')
                fout.flush()
            slines = []

        slines.append(sline.strip())
        count += 1
    if slines != []:
        hypotheses_batch = bart.sample(slines, beam=1, max_len_b=20)
        for hypothesis in hypotheses_batch:
            fout.write(hypothesis + '\n')
            fout.flush()