pipeline/translate/extract_best.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import absolute_import, division, print_function

import argparse
import collections
import math
import re
import sys


def main():
    args = parse_args()

    if args.metric == "bleu":
        score_function = compute_bleu
    elif args.metric == "sacrebleu":
        global sacrebleu
        import sacrebleu

        score_function = compute_sacrebleu
    elif args.metric == "chrf":
        global sacrebleu
        import sacrebleu

        score_function = compute_chrf
    else:
        sys.stderr.write("Unrecognized metric: {}\n".format(args.metric))
        pass

    if args.toolkit == "marian":
        marian_best_bleu(args, score_function)
    elif args.toolkit == "t2t":
        t2t_best_bleu(args, score_function)
        pass


def t2t_best_bleu(args, score_function):
    for i, ref_line in enumerate(args.references):
        refs = ref_line.strip().split("\n")
        if args.debpe:
            refs = [re.sub(r"@@ +", "", r) for r in refs]
            pass
        texts = next(args.nbest).strip().split("\t")
        if args.debpe:
            texts = [re.sub(r"@@ +", "", t) for t in texts]
            pass
        refs = [r.split() for r in refs]
        scores = [score_function(refs, t.split()) for t in texts]
        best_txt = texts[scores.index(max(scores))]

        args.output.write("{}\n".format(best_txt))
        if args.debug:
            sys.stderr.write("{}: {}\n".format(i, scores))
            pass
        if i % 100000 == 0 and i > 0:
            sys.stderr.write("[{}]\n".format(i))
            pass
        pass


def marian_best_bleu(args, score_function):
    prev_line = None
    for i, ref_line in enumerate(args.references):
        refs = ref_line.strip().split("\n")
        if args.debpe:
            refs = [re.sub(r"@@ +", "", r) for r in refs]

        texts = []
        while True:
            if prev_line:
                # CTranslate2 can output empty text, for example:
                # 10181 ||| .GDFMAKERPROJECTファイルを開くには?
                # 10181 ||| .GDMAKERPROJECTファイルを開くには?
                # 10181 ||| .GDFMAKERPROJECTファイルを開くには?
                # 10181 ||| .GDFMakerPROJECTファイルを開くには?
                # 10181 ||| .GDFAKERPROJECTファイルを開くには?
                # 10181 ||| .GDMakerPROJECTファイルを開くには?
                # 10181 ||| .GDFMAKERPROJECTファイルを開くには。
                # 10181 |||
                # Marian also outputs scores, for example:
                # 0 ||| Реформа, направленная на выдвижение условий, идет слишком медленно. ||| F0= -9.21191 F1= -11.53 ||| -1.22059
                fields = prev_line.rstrip("\n").split(" ||| ")
                if len(fields) == 1:
                    # handle "10181 |||"
                    fields = fields[0].split()[0], ""

                idx = int(fields[0])
                if idx == i:
                    texts.append(fields[1])
                else:
                    break

            prev_line = next(args.nbest, None)
            if not prev_line:
                break

        if args.debpe:
            texts = [re.sub(r"@@ +", "", t) for t in texts]
        refs = [r.split() for r in refs]
        scores = [score_function(refs, t.split()) for t in texts]
        best_txt = texts[scores.index(max(scores))]

        args.output.write("{}\n".format(best_txt))
        if args.debug:
            sys.stderr.write("{}: {}\n".format(i, scores))

        if i % 100000 == 0 and i > 0:
            sys.stderr.write("[{}]\n".format(i))


def compute_chrf(references, translation):
    hypo = " ".join(translation)
    refs = [" ".join(r) for r in references]
    return sacrebleu.sentence_chrf(hypo, refs).score


def compute_sacrebleu(references, translation):
    hypo = " ".join(translation)
    refs = [" ".join(r) for r in references]
    return sacrebleu.sentence_bleu(hypo, refs).score


def compute_bleu(references, translation, max_order=4):
    precisions = get_ngram_precisions(references, translation, max_order)
    if min(precisions) > 0:
        p_log_sum = sum((1.0 / max_order) * math.log(p) for p in precisions)
        geo_mean = math.exp(p_log_sum)
    else:
        geo_mean = 0

    bp = get_brevity_penalty(references, translation)
    return geo_mean * bp


def get_brevity_penalty(references, translation):
    reference_length = min(len(r) for r in references)
    translation_length = len(translation)
    ratio = float(translation_length) / reference_length
    if ratio > 1.0 or ratio == 0.0:
        bp = 1.0
    else:
        bp = math.exp(1 - 1.0 / ratio)
    return bp


def get_ngram_precisions(references, translation, max_order=4):
    matches_by_order = [0] * max_order
    possible_matches_by_order = [0] * max_order

    merged_ref_ngram_counts = collections.Counter()
    for reference in references:
        merged_ref_ngram_counts |= get_ngrams(reference, max_order)
    translation_ngram_counts = get_ngrams(translation, max_order)
    overlap = translation_ngram_counts & merged_ref_ngram_counts
    for ngram in overlap:
        matches_by_order[len(ngram) - 1] += overlap[ngram]
    for order in range(1, max_order + 1):
        possible_matches = len(translation) - order + 1
        if possible_matches > 0:
            possible_matches_by_order[order - 1] += possible_matches

    precisions = [0] * max_order
    for i in range(0, max_order):
        # smoothing
        if matches_by_order[i] == 0 and possible_matches_by_order[i] == 0:
            precisions[i] = 0.0
        else:
            precisions[i] = (matches_by_order[i] + 1.0) / (possible_matches_by_order[i] + 1.0)
    return precisions


def get_ngrams(segment, max_order):
    ngram_counts = collections.Counter()
    for order in range(1, max_order + 1):
        for i in range(0, len(segment) - order + 1):
            ngram = tuple(segment[i : i + order])
            ngram_counts[ngram] += 1
    return ngram_counts


def parse_args():
    from argparse import FileType

    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--nbest", type=FileType("r"), default=sys.stdin)
    parser.add_argument("-r", "--references", type=FileType("r"), required=True)
    parser.add_argument("-o", "--output", type=FileType("w"), default=sys.stdout)
    parser.add_argument("-m", "--metric", default="bleu")
    parser.add_argument("--debpe", action="store_true")
    parser.add_argument("-d", "--debug", action="store_true")
    parser.add_argument("-t", "--toolkit", default="marian", help="Toolkit: 'marian' or 't2t'")
    return parser.parse_args()


if __name__ == "__main__":
    main()