chore: import upstream snapshot with attribution

2026-07-13 13:24:13 +08:00
commit 1037506f2e
6050 changed files with 1731598 additions and 0 deletions
@@ -0,0 +1,24 @@
+# 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.
+
+import importlib
+import os
+
+from fairseq import registry
+
+
+build_agent, register_agent, MONOTONIC_AGENT, _ = registry.setup_registry(
+    "--agent-type"
+)
+
+
+DEFAULT_EOS = "</s>"
+GET = 0
+SEND = 1
+
+for file in os.listdir(os.path.dirname(__file__)):
+    if file.endswith(".py") and not file.startswith("_"):
+        module = file[: file.find(".py")]
+        importlib.import_module("agents." + module)
@@ -0,0 +1,67 @@
+# 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.
+
+import time
+from functools import partial
+from multiprocessing.pool import ThreadPool as Pool
+
+from . import DEFAULT_EOS, GET, SEND
+
+
+class Agent(object):
+    "an agent needs to follow this pattern"
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def init_states(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def update_states(self, states, new_state):
+        raise NotImplementedError
+
+    def finish_eval(self, states, new_state):
+        raise NotImplementedError
+
+    def policy(self, state):
+        raise NotImplementedError
+
+    def reset(self):
+        raise NotImplementedError
+
+    def decode(self, session, low=0, high=100000, num_thread=10):
+        corpus_info = session.corpus_info()
+        high = min(corpus_info["num_sentences"] - 1, high)
+        if low >= high:
+            return
+
+        t0 = time.time()
+        if num_thread > 1:
+            with Pool(10) as p:
+                p.map(
+                    partial(self._decode_one, session),
+                    [sent_id for sent_id in range(low, high + 1)],
+                )
+        else:
+            for sent_id in range(low, high + 1):
+                self._decode_one(session, sent_id)
+
+        print(f"Finished {low} to {high} in {time.time() - t0}s")
+
+    def _decode_one(self, session, sent_id):
+        action = {}
+        self.reset()
+        states = self.init_states()
+        while action.get("value", None) != DEFAULT_EOS:
+            # take an action
+            action = self.policy(states)
+
+            if action["key"] == GET:
+                new_states = session.get_src(sent_id, action["value"])
+                states = self.update_states(states, new_states)
+
+            elif action["key"] == SEND:
+                session.send_hypo(sent_id, action["value"])
+        print(" ".join(states["tokens"]["tgt"]))
@@ -0,0 +1,167 @@
+# 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.
+
+import json
+import os
+
+from fairseq import checkpoint_utils, tasks, utils
+
+from . import DEFAULT_EOS, GET, SEND
+from .agent import Agent
+
+
+class SimulTransAgent(Agent):
+    def __init__(self, args):
+        # Load Model
+        self.load_model(args)
+
+        # build word spliter
+        self.build_word_splitter(args)
+
+        self.max_len = args.max_len
+
+        self.eos = DEFAULT_EOS
+
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument('--model-path', type=str, required=True,
+                            help='path to your pretrained model.')
+        parser.add_argument("--data-bin", type=str, required=True,
+                            help="Path of data binary")
+        parser.add_argument("--user-dir", type=str, default="example/simultaneous_translation",
+                            help="User directory for simultaneous translation")
+        parser.add_argument("--src-splitter-type", type=str, default=None,
+                            help="Subword splitter type for source text")
+        parser.add_argument("--tgt-splitter-type", type=str, default=None,
+                            help="Subword splitter type for target text")
+        parser.add_argument("--src-splitter-path", type=str, default=None,
+                            help="Subword splitter model path for source text")
+        parser.add_argument("--tgt-splitter-path", type=str, default=None,
+                            help="Subword splitter model path for target text")
+        parser.add_argument("--max-len", type=int, default=150,
+                            help="Maximum length difference between source and target prediction")
+        parser.add_argument('--model-overrides', default="{}", type=str, metavar='DICT',
+                            help='A dictionary used to override model args at generation '
+                                 'that were used during model training')
+        # fmt: on
+        return parser
+
+    def load_dictionary(self, task):
+        raise NotImplementedError
+
+    def load_model(self, args):
+        args.user_dir = os.path.join(os.path.dirname(__file__), "..", "..")
+        utils.import_user_module(args)
+        filename = args.model_path
+        if not os.path.exists(filename):
+            raise IOError("Model file not found: {}".format(filename))
+
+        state = checkpoint_utils.load_checkpoint_to_cpu(
+            filename, json.loads(args.model_overrides)
+        )
+
+        saved_args = state["args"]
+        saved_args.data = args.data_bin
+
+        task = tasks.setup_task(saved_args)
+
+        # build model for ensemble
+        self.model = task.build_model(saved_args)
+        self.model.load_state_dict(state["model"], strict=True)
+
+        # Set dictionary
+        self.load_dictionary(task)
+
+    def init_states(self):
+        return {
+            "indices": {"src": [], "tgt": []},
+            "tokens": {"src": [], "tgt": []},
+            "segments": {"src": [], "tgt": []},
+            "steps": {"src": 0, "tgt": 0},
+            "finished": False,
+            "finish_read": False,
+            "model_states": {},
+        }
+
+    def update_states(self, states, new_state):
+        raise NotImplementedError
+
+    def policy(self, states):
+        # Read and Write policy
+        action = None
+
+        while action is None:
+            if states["finished"]:
+                # Finish the hypo by sending eos to server
+                return self.finish_action()
+
+            # Model make decision given current states
+            decision = self.model.decision_from_states(states)
+
+            if decision == 0 and not self.finish_read(states):
+                # READ
+                action = self.read_action(states)
+            else:
+                # WRITE
+                action = self.write_action(states)
+
+            # None means we make decision again but not sending server anything
+            # This happened when read a bufffered token
+            # Or predict a subword
+        return action
+
+    def finish_read(self, states):
+        raise NotImplementedError
+
+    def write_action(self, states):
+        token, index = self.model.predict_from_states(states)
+
+        if (
+            index == self.dict["tgt"].eos()
+            or len(states["tokens"]["tgt"]) > self.max_len
+        ):
+            # Finish this sentence is predict EOS
+            states["finished"] = True
+            end_idx_last_full_word = self._target_length(states)
+
+        else:
+            states["tokens"]["tgt"] += [token]
+            end_idx_last_full_word = self.word_splitter["tgt"].end_idx_last_full_word(
+                states["tokens"]["tgt"]
+            )
+            self._append_indices(states, [index], "tgt")
+
+        if end_idx_last_full_word > states["steps"]["tgt"]:
+            # Only sent detokenized full words to the server
+            word = self.word_splitter["tgt"].merge(
+                states["tokens"]["tgt"][states["steps"]["tgt"] : end_idx_last_full_word]
+            )
+            states["steps"]["tgt"] = end_idx_last_full_word
+            states["segments"]["tgt"] += [word]
+
+            return {"key": SEND, "value": word}
+        else:
+            return None
+
+    def read_action(self, states):
+        return {"key": GET, "value": None}
+
+    def finish_action(self):
+        return {"key": SEND, "value": DEFAULT_EOS}
+
+    def reset(self):
+        pass
+
+    def finish_eval(self, states, new_state):
+        if len(new_state) == 0 and len(states["indices"]["src"]) == 0:
+            return True
+        return False
+
+    def _append_indices(self, states, new_indices, key):
+        states["indices"][key] += new_indices
+
+    def _target_length(self, states):
+        return len(states["tokens"]["tgt"])
@@ -0,0 +1,81 @@
+# 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.
+
+from . import DEFAULT_EOS, GET, register_agent
+from .simul_trans_agent import SimulTransAgent
+from .word_splitter import SPLITTER_DICT
+
+
+@register_agent("simul_trans_text")
+class SimulTransTextAgent(SimulTransAgent):
+    def build_word_splitter(self, args):
+        self.word_splitter = {}
+
+        self.word_splitter["src"] = SPLITTER_DICT[args.src_splitter_type](
+            getattr(args, f"src_splitter_path")
+        )
+        self.word_splitter["tgt"] = SPLITTER_DICT[args.tgt_splitter_type](
+            getattr(args, f"tgt_splitter_path")
+        )
+
+    def load_dictionary(self, task):
+        self.dict = {}
+        self.dict["tgt"] = task.target_dictionary
+        self.dict["src"] = task.source_dictionary
+
+    def update_states(self, states, new_state):
+        if states["finish_read"]:
+            return states
+
+        new_word = new_state["segment"]
+
+        # Split words and index the token
+        if new_word not in [DEFAULT_EOS]:
+            tokens = self.word_splitter["src"].split(new_word)
+            # Get indices from dictionary
+            # You can change to you own dictionary
+            indices = (
+                self.dict["src"]
+                .encode_line(
+                    tokens,
+                    line_tokenizer=lambda x: x,
+                    add_if_not_exist=False,
+                    append_eos=False,
+                )
+                .tolist()
+            )
+        else:
+            tokens = [new_word]
+            indices = [self.dict["src"].eos()]
+            states["finish_read"] = True
+
+        # Update states
+        states["segments"]["src"] += [new_word]
+        states["tokens"]["src"] += tokens
+        self._append_indices(states, indices, "src")
+
+        return states
+
+    def read_action(self, states):
+        # Increase source step by one
+        states["steps"]["src"] += 1
+
+        # At leat one word is read
+        if len(states["tokens"]["src"]) == 0:
+            return {"key": GET, "value": None}
+
+        # Only request new word if there is no buffered tokens
+        if len(states["tokens"]["src"]) <= states["steps"]["src"]:
+            return {"key": GET, "value": None}
+
+        return None
+
+    def finish_read(self, states):
+        # The first means all segments (full words) has been read from server
+        # The second means all tokens (subwords) has been read locally
+        return (
+            states["finish_read"]
+            and len(states["tokens"]["src"]) == states["steps"]["src"]
+        )
@@ -0,0 +1,91 @@
+# 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.
+
+
+class SubwordSplitter(object):
+    def process_line(self, string):
+        raise NotImplementedError
+
+    def split(self, string):
+        raise NotImplementedError
+
+
+class NoneWordSplitter(object):
+    def __init__(self, model):
+        pass
+
+    def split(self, string):
+        return [string]
+
+    def process_line(self, string):
+        return [string]
+
+    def finished_word(self, string):
+        return True
+
+    def merge(self, list_of_string):
+        return "".join(list_of_string)
+
+    def last_full_word_step(self, tokens, step):
+        return len(tokens)
+
+    def end_idx_last_full_word(self, tokens):
+        return len(tokens)
+
+
+class BPEWordSplitter(object):
+    # TODO: lock back here
+    def __init__(self, model_path):
+        super().__init__()
+        from subword_nmt.apply_bpe import BPE
+
+        with open(model_path) as f:
+            self.model = BPE(f)
+
+    def split(self, string):
+        return self.model.process_line(string).split()
+
+    def end_idx_last_full_word(self, tokens):
+        # Begin of word indices
+        bow_indices = [0] + [i + 1 for i, t in enumerate(tokens[1:]) if t[-2:] != "@@"]
+
+        if len(bow_indices) < 2:
+            return 0
+        else:
+            return bow_indices[-1]
+
+    def merge(self, list_of_string):
+        return " ".join([item.replace("@@", "") for item in list_of_string])
+
+
+class SentencePieceModelWordSplitter(object):
+    def __init__(self, model_path):
+        super().__init__()
+        import sentencepiece as spm
+
+        self.model = spm.SentencePieceProcessor()
+        self.model.Load(model_path)
+
+    def split(self, string):
+        return self.model.EncodeAsPieces(string)
+
+    def end_idx_last_full_word(self, tokens):
+        # Begin of word indices
+        bow_indices = [i for i, t in enumerate(tokens) if t[0] == "\u2581"]
+
+        if len(bow_indices) < 2:
+            return 0
+        else:
+            return bow_indices[-1]
+
+    def merge(self, list_of_string):
+        return self.model.DecodePieces(list_of_string)
+
+
+SPLITTER_DICT = {
+    None: NoneWordSplitter,
+    "BPE": BPEWordSplitter,
+    "SentencePieceModel": SentencePieceModelWordSplitter,
+}