app/ai4bharat/transliteration/transformer/custom_interactive.py

#!/usr/bin/env python3 -u
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
"""
Translate raw text with a trained model. Batches data on-the-fly.
"""

import ast
import fileinput
import logging
import math
import os
import sys
import time
from argparse import Namespace
from collections import namedtuple

import numpy as np
import torch

from fairseq import checkpoint_utils, distributed_utils, options, tasks, utils
from fairseq.dataclass.configs import FairseqConfig
from fairseq.dataclass.utils import convert_namespace_to_omegaconf
from fairseq.token_generation_constraints import pack_constraints, unpack_constraints
from fairseq_cli.generate import get_symbols_to_strip_from_output

Batch = namedtuple("Batch", "ids src_tokens src_lengths constraints")
Translation = namedtuple("Translation", "src_str hypos pos_scores alignments")

def make_batches(lines, cfg, task, max_positions, encode_fn):
    def encode_fn_target(x):
        return encode_fn(x)


    if cfg.generation.constraints:
        # Strip (tab-delimited) contraints, if present, from input lines,
        # store them in batch_constraints
        batch_constraints = [list() for _ in lines]
        for i, line in enumerate(lines):
            if "\t" in line:
                lines[i], *batch_constraints[i] = line.split("\t")

        # Convert each List[str] to List[Tensor]
        for i, constraint_list in enumerate(batch_constraints):
            batch_constraints[i] = [
                task.target_dictionary.encode_line(
                    encode_fn_target(constraint),
                    append_eos=False,
                    add_if_not_exist=False,
                )
                for constraint in constraint_list
            ]

    if cfg.generation.constraints:
        constraints_tensor = pack_constraints(batch_constraints)
    else:
        constraints_tensor = None

    tokens, lengths = task.get_interactive_tokens_and_lengths(lines, encode_fn)

    itr = task.get_batch_iterator(
        dataset=task.build_dataset_for_inference(
            tokens, lengths, constraints=constraints_tensor
        ),
        max_tokens=cfg.dataset.max_tokens,
        max_sentences=cfg.dataset.batch_size,
        max_positions=max_positions,
        ignore_invalid_inputs=cfg.dataset.skip_invalid_size_inputs_valid_test,
    ).next_epoch_itr(shuffle=False)
    for batch in itr:
        ids = batch["id"]
        src_tokens = batch["net_input"]["src_tokens"]
        src_lengths = batch["net_input"]["src_lengths"]
        constraints = batch.get("constraints", None)

        yield Batch(
            ids=ids,
            src_tokens=src_tokens,
            src_lengths=src_lengths,
            constraints=constraints,
        )


# added
class Transliterator:
    def __init__(
        self, data_bin_dir, model_checkpoint_path, lang_pairs_csv, lang_list_file, beam, batch_size = 32,
    ):

        self.parser = options.get_interactive_generation_parser()

        # buffer_size is currently not used but we just initialize it to batch
        # size + 1 to avoid any assertion errors.
        
        self.parser.set_defaults(
            path = model_checkpoint_path,
            num_wokers = -1,
            batch_size = batch_size,
            buffer_size = batch_size + 1,
            task = "translation_multi_simple_epoch",
            beam = beam,
            # nbest = nbest,
            # source_lang = 'en' ,
            # target_lang = 'mlt' ,
            # encoder_langtok = "tgt" ,
            # lang_dict = "lang_list.txt"
        )
        
        self.args = options.parse_args_and_arch(self.parser, input_args = [data_bin_dir] )
        
        self.args.skip_invalid_size_inputs_valid_test = False
        # self.args.lang_pairs = "en-as,en-bn,en-gom,en-gu,en-hi,en-kn,en-ks,en-mai,en-ml,en-mr,en-ne,en-or,en-pa,en-sa,en-sd,en-si,en-ta,en-te,en-ur"
        self.args.lang_pairs = lang_pairs_csv
        # self.args.source_lang = 'en'
        # self.args.target_lang = 'bn'
        # self.args.encoder_langtok = 'tgt'
        self.args.lang_dict = lang_list_file

        self.cfg = convert_namespace_to_omegaconf(self.args)

        if isinstance(self.cfg, Namespace):
            self.cfg = convert_namespace_to_omegaconf(self.cfg)

        # start_time = time.time()
        self.total_translate_time = 0

        utils.import_user_module(self.cfg.common)

        if self.cfg.interactive.buffer_size < 1:
            self.cfg.interactive.buffer_size = 1
        if self.cfg.dataset.max_tokens is None and self.cfg.dataset.batch_size is None:
            self.cfg.dataset.batch_size = 1

        assert (
            not self.cfg.generation.sampling or self.cfg.generation.nbest == self.cfg.generation.beam
        ), "--sampling requires --nbest to be equal to --beam"
        assert (
            not self.cfg.dataset.batch_size
            or self.cfg.dataset.batch_size <= self.cfg.interactive.buffer_size
        ), "--batch-size cannot be larger than --buffer-size"


        # FIXME: Following lines are commented out by GokulNC due to some erros on CPU
        # GokulNC: Why do we need seeding for inference? What is stochastic decoding?

        # # Fix seed for stochastic decoding
        # if self.cfg.common.seed is not None and not self.cfg.generation.no_seed_provided:
        #     np.random.seed(self.cfg.common.seed)
        #     utils.set_torch_seed(self.cfg.common.seed)

        self.use_cuda = torch.cuda.is_available() and not self.cfg.common.cpu

        # Setup task, e.g., translation
        self.task = tasks.setup_task(self.cfg.task)

        # Load ensemble
        overrides = ast.literal_eval(self.cfg.common_eval.model_overrides)
        # logger.info("loading model(s) from {}".format(self.cfg.common_eval.path))
        self.models, _model_args = checkpoint_utils.load_model_ensemble(
            utils.split_paths(self.cfg.common_eval.path),
            arg_overrides=overrides,
            task=self.task,
            suffix=self.cfg.checkpoint.checkpoint_suffix,
            strict=(self.cfg.checkpoint.checkpoint_shard_count == 1),
            num_shards=self.cfg.checkpoint.checkpoint_shard_count,
        )

        # Set dictionaries
        self.src_dict = self.task.source_dictionary
        self.tgt_dict = self.task.target_dictionary
        
        # print("src dict",self.src_dict)
        # print("src dict",self.tgt_dict)
        # print("self.src_dict.__len__() : ",self.src_dict.symbols)
        # print("self.src_dict.__len__() : ",self.tgt_dict.symbols)

        # self.cfg.common.fp16 = True
        # Optimize ensemble for generation
        for i in range(len(self.models)):
            if self.models[i] is None:
                continue
            if self.cfg.common.fp16:
                self.models[i].half()
            
            if self.use_cuda and not self.cfg.distributed_training.pipeline_model_parallel:
                self.models[i].cuda()
            self.models[i].prepare_for_inference_(self.cfg)
            # # Quantize
            # self.models[i] = torch.quantization.quantize_dynamic(
            #     self.models[i], {torch.nn.Linear}, dtype=torch.qint8
            # )
            # # Torchscript
            # self.models[i] = torch.jit.script(self.models[i])

        # Initialize generator
        self.generator = self.task.build_generator(self.models, self.cfg.generation)

        # Handle tokenization and BPE
        self.tokenizer = self.task.build_tokenizer(self.cfg.tokenizer)
        self.bpe = self.task.build_bpe(self.cfg.bpe)

        # Load alignment dictionary for unknown word replacement
        # (None if no unknown word replacement, empty if no path to align dictionary)
        self.align_dict = utils.load_align_dict(self.cfg.generation.replace_unk)

        self.max_positions = utils.resolve_max_positions(
            self.task.max_positions(), *[model.max_positions() for model in self.models]
        )

    def encode_fn(self, x):
        if self.tokenizer is not None:
            x = self.tokenizer.encode(x)
        if self.bpe is not None:
            x = self.bpe.encode(x)
        return x

    def decode_fn(self, x):
        if self.bpe is not None:
            x = self.bpe.decode(x)
        if self.tokenizer is not None:
            x = self.tokenizer.decode(x)
        return x

    def translate(self, inputs, nbest=1):

        start_id = 0
        # for inputs in buffered_read(self.cfg.interactive.input, self.cfg.interactive.buffer_size):

        results = []
        for batch in make_batches(inputs, self.cfg, self.task, self.max_positions, self.encode_fn):
            bsz = batch.src_tokens.size(0)
            src_tokens = batch.src_tokens
            src_lengths = batch.src_lengths
            constraints = batch.constraints
            if self.use_cuda:
                src_tokens = src_tokens.cuda()
                src_lengths = src_lengths.cuda()
                if constraints is not None:
                    constraints = constraints.cuda()

            sample = {
                "net_input": {
                    "src_tokens": src_tokens,
                    "src_lengths": src_lengths,
                },
            }

            translate_start_time = time.time()
            translations = self.task.inference_step(
                self.generator, self.models, sample, constraints=constraints
            )
            translate_time = time.time() - translate_start_time
            self.total_translate_time += translate_time
            list_constraints = [[] for _ in range(bsz)]
            if self.cfg.generation.constraints:
                list_constraints = [unpack_constraints(c) for c in constraints]
            for i, (id, hypos) in enumerate(zip(batch.ids.tolist(), translations)):
                src_tokens_i = utils.strip_pad(src_tokens[i], self.tgt_dict.pad())
                constraints = list_constraints[i]
                results.append(
                    (
                        start_id + id,
                        src_tokens_i,
                        hypos,
                        {
                            "constraints": constraints,
                            "time": translate_time / len(translations),
                        },
                    )
                )

        # sort output to match input order
        result_str = ""
        for id_, src_tokens, hypos, info in sorted(results, key=lambda x: x[0]):

            src_str = ""
            if self.src_dict is not None:
                src_str = self.src_dict.string(src_tokens, self.cfg.common_eval.post_process)

                
                # print("S-{}\t{}".format(id_, src_str))
                result_str += "S-{}\t{}".format(id_, src_str) + '\n'

                # print("W-{}\t{:.3f}\tseconds".format(id_, info["time"]))
                result_str += "W-{}\t{:.3f}\tseconds".format(id_, info["time"]) + '\n'

                for constraint in info["constraints"]:
                    # print(
                    #     "C-{}\t{}".format(
                    #         id_,
                    #         self.tgt_dict.string(constraint, self.cfg.common_eval.post_process),
                    #     )
                    # )
                    result_str += "C-{}\t{}".format(
                            id_,
                            self.tgt_dict.string(constraint, self.cfg.common_eval.post_process),
                        ) + '\n'

            # Process top predictions
            for hypo in hypos[: min(len(hypos), nbest)]:
                hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
                    hypo_tokens=hypo["tokens"].int().cpu(),
                    src_str=src_str,
                    alignment=hypo["alignment"],
                    align_dict=self.align_dict,
                    tgt_dict=self.tgt_dict,
                    remove_bpe=self.cfg.common_eval.post_process,
                    extra_symbols_to_ignore=get_symbols_to_strip_from_output(self.generator),
                )
                detok_hypo_str = self.decode_fn(hypo_str)
                score = hypo["score"] / math.log(2)  # convert to base 2
                # original hypothesis (after tokenization and BPE)
                # print("H-{}\t{}\t{}".format(id_, score, hypo_str))
                result_str += "H-{}\t{}\t{}".format(id_, score, hypo_str) + '\n'

                # detokenized hypothesis
                # print("D-{}\t{}\t{}".format(id_, score, detok_hypo_str))
                result_str += "D-{}\t{}\t{}".format(id_, score, detok_hypo_str) + '\n'
                
                # print(
                #     "P-{}\t{}".format(
                #         id_,
                #         " ".join(
                #             map(
                #                 lambda x: "{:.4f}".format(x),
                #                 # convert from base e to base 2
                #                 hypo["positional_scores"].div_(math.log(2)).tolist(),
                #             )
                #         ),
                #     )
                # )
                result_str += "P-{}\t{}".format(
                        id_,
                        " ".join(
                            map(
                                lambda x: "{:.4f}".format(x),
                                # convert from base e to base 2
                                hypo["positional_scores"].div_(math.log(2)).tolist(),
                            )
                        ),
                    ) + '\n'

                if self.cfg.generation.print_alignment:
                    alignment_str = " ".join(
                        ["{}-{}".format(src, tgt) for src, tgt in alignment]
                    )
                    # print("A-{}\t{}".format(id_, alignment_str))
                    result_str += "A-{}\t{}".format(id_, alignment_str) + '\n'

            # # update running id_ counter
            # start_id += len(inputs)
        return result_str