<commit_before>import multiprocessing
import os
import subprocess

import scripts
from moses import MosesFeature, Moses
from scripts.libs import fileutils, shell
from scripts.mt import ParallelCorpus

__author__ = '<NAME>'


class WordAligner:
    available_types = ['FastAlign', 'mgizapp']

    def __init__(self):
        pass

    @staticmethod
    def instantiate(type_name):
        if type_name == 'FastAlign':
            return FastAlign()
        elif type_name == 'mgizapp':
            return MGizaPP()
        else:
            raise NameError('Invalid Word Aligner type: ' + str(type_name))

    def align(self, corpus, langs, model_dir, working_dir='.', log_file=None):
        if not os.path.isdir(working_dir):
            fileutils.makedirs(working_dir, exist_ok=True)


class SuffixArraysPhraseTable(MosesFeature):
    injector_section = 'suffixarrays'
    injectable_fields = {
        'sample': ('number of samples for phrase table', int, 1000),
        'method': ('sampling method for suffix array phrase table', str, 'ranked3'),
    }

    def __init__(self, model, langs):
        MosesFeature.__init__(self, 'Mmsapt')

        self._sample = None  # Injected
        self._method = None  # Injected

        self._model = model
        self._source_lang = langs[0]
        self._target_lang = langs[1]

        self._symal_bin = os.path.join(Moses.bin_path, 'bin', 'symal')
        self._symal2mam_bin = os.path.join(Moses.bin_path, 'bin', 'symal2mam')
        self._mttbuild_bin = os.path.join(Moses.bin_path, 'bin', 'mtt-build')
        self._mmlexbuild_bin = os.path.join(Moses.bin_path, 'bin', 'mmlex-build')

    def _get_model_basename(self):
        return os.path.join(self._model, 'model')

    def get_iniline(self):
        template = 'path={model}. L1={source_lang} L2={target_lang} output-factor=0 sample={sample} method={method} ' \
                   'pfwd=g pbwd=g logcnt=0 coh=0 prov=0 rare=0 unal=0 smooth=.01 lexalpha=0 lr-func=DM0 bias-loglevel=1'

        return template.format(model=self.get_relpath(self._get_model_basename()),
                               source_lang=self._source_lang, target_lang=self._target_lang,
                               sample=self._sample, method=self._method)

    def train(self, corpora, aligner, working_dir='.', log_file=None):
        if os.path.isdir(self._model) and len(os.listdir(self._model)) > 0:
            raise Exception('Model already exists at ' + self._model)

        if not os.path.isdir(self._model):
            fileutils.makedirs(self._model, exist_ok=True)

        if not os.path.isdir(working_dir):
            fileutils.makedirs(working_dir, exist_ok=True)

        l1 = self._source_lang
        l2 = self._target_lang
        langs = (l1, l2)
        langs_suffix = l1 + '-' + l2

        mct_base = self._get_model_basename()
        dmp_file = mct_base + '.dmp'
        mam_file = mct_base + '.' + langs_suffix + '.mam'
        lex_file = mct_base + '.' + langs_suffix + '.lex'

        log = shell.DEVNULL

        try:
            if log_file is not None:
                log = open(log_file, 'a')

            merged_corpus = ParallelCorpus(os.path.basename(mct_base), working_dir, langs)

            # Create merged corpus and domains list file (dmp)
            fileutils.merge([corpus.get_file(l1) for corpus in corpora], merged_corpus.get_file(l1))
            fileutils.merge([corpus.get_file(l2) for corpus in corpora], merged_corpus.get_file(l2))
            with open(dmp_file, 'w') as dmp:
                for corpus in corpora:
                    dmp.write(str(corpus.name) + ' ' + str(corpus.count_lines()) + '\n')

            # Create alignments in 'bal' file and symmetrize
            bal_file = aligner.align(merged_corpus, langs, self._model, working_dir, log_file)

            symal_file = os.path.join(working_dir, 'alignments.' + langs_suffix + '.symal')
            symal_command = [self._symal_bin, '-a=g', '-d=yes', '-f=yes', '-b=yes']
            with open(bal_file) as stdin:
                with open(symal_file, 'w') as stdout:
                    shell.execute(symal_command, stdin=stdin, stdout=stdout, stderr=log)

            # Execute mtt-build
            mttbuild_command = self._get_mttbuild_command(mct_base, dmp_file, l1)
            with open(merged_corpus.get_file(l1)) as stdin:
                shell.execute(mttbuild_command, stdin=stdin, stdout=log, stderr=log)

            mttbuild_command = self._get_mttbuild_command(mct_base, dmp_file, l2)
            with open(merged_corpus.get_file(l2)) as stdin:
                shell.execute(mttbuild_command, stdin=stdin, stdout=log, stderr=log)

            # Create 'mam' file
            mam_command = [self._symal2mam_bin, mam_file]
            with open(symal_file) as stdin:
                shell.execute(mam_command, stdin=stdin, stdout=log, stderr=log)

            # Create 'lex' file
            lex_command = [self._mmlexbuild_bin, mct_base + '.', l1, l2, '-o', lex_file]
            shell.execute(lex_command, stdout=log, stderr=log)
        finally:
            if log_file is not None:
                log.close()

    def _get_mttbuild_command(self, mct_base, dmp_file, lang):
        output = mct_base + '.' + lang
        return [self._mttbuild_bin, '-i', '-o', output, '-m', dmp_file, '-g']


class FastAlign(WordAligner):
    def __init__(self):
        WordAligner.__init__(self)

        self._align_bin = os.path.join(scripts.BIN_DIR, 'fastalign-maurobuild', 'fast_align')

    def align(self, corpus, langs, model_dir, working_dir='.', log_file=None):
        WordAligner.align(self, corpus, langs, working_dir, log_file)

        l1 = langs[0]
        l2 = langs[1]
        corpus_name = 'corpus'
        langs_suffix = l1 + '-' + l2

        fwd_file = os.path.join(working_dir, corpus_name + '.' + langs_suffix + '.fwd')
        bwd_file = os.path.join(working_dir, corpus_name + '.' + langs_suffix + '.bwd')
        bal_file = os.path.join(working_dir, corpus_name + '.' + langs_suffix + '.bal')
        aligned_file_path = os.path.join(working_dir, corpus_name + '.' + langs_suffix + '.aligned')

        corpus_l1 = corpus.get_file(l1)
        corpus_l2 = corpus.get_file(l2)

        log = shell.DEVNULL

        try:
            if log_file is not None:
                log = open(log_file, 'a')

            with open(corpus_l1) as source_corpus, open(corpus_l2) as target_corpus, open(aligned_file_path, 'w') as aligned_file:
                for x, y in zip(source_corpus, target_corpus):
                    aligned_file.write(x.strip() + " ||| " + y.strip() + "\n")

            # Forward alignments
            fwd_model = os.path.join(model_dir,  'model.align.fwd')
            command = [self._align_bin, '-d', '-v', '-o', '-B', '-p', fwd_model, '-i', aligned_file_path]
            with open(fwd_file, 'w') as stdout:
                shell.execute(command, stdout=stdout, stderr=log)

            # Forward alignments
            bwd_model = os.path.join(model_dir,  'model.align.bwd')
            command = [self._align_bin, '-d', '-v', '-o', '-B','-p', bwd_model, '-r', '-i', aligned_file_path]
            with open(bwd_file, 'w') as stdout:
                shell.execute(command, stdout=stdout, stderr=log)


        finally:
            if log_file is not None:
                log.close()

        encoder = _FastAlignBALEncoder(corpus, langs, fwd_file, bwd_file)
        encoder.encode(bal_file)

        return bal_file


class _FastAlignBALEncoder:
    def __init__(self, corpus, langs, fwd, bwd):
        self._corpus_l1 = corpus.get_file(langs[0])
        self._corpus_l2 = corpus.get_file(langs[1])
        self._fwd = fwd
        self._bwd = bwd

    def encode(self, output):
        files = []

        try:
            files.append(open(self._corpus_l1))
            files.append(open(self._corpus_l2))
            files.append(open(self._fwd))
            files.append(open(self._bwd))
            files.append(open(output, 'w'))

            c1, c2, fwd, bwd, bal = files

            for t1 in c1:
                t1 = t1.strip().split()<commit_msg>Added NumberOfThreads option to fast_align during training to use all available CPUs<commit_after>import multiprocessing
import os

import scripts
from moses import MosesFeature, Moses
from scripts.libs import fileutils, shell
from scripts.mt import ParallelCorpus

__author__ = '<NAME>'


class WordAligner:
    available_types = ['FastAlign', 'mgizapp']

    def __init__(self):
        pass

    @staticmethod
    def instantiate(type_name):
        if type_name == 'FastAlign':
            return FastAlign()
        elif type_name == 'mgizapp':
            return MGizaPP()
        else:
            raise NameError('Invalid Word Aligner type: ' + str(type_name))

    def align(self, corpus, langs, model_dir, working_dir='.', log_file=None):
        if not os.path.isdir(working_dir):
            fileutils.makedirs(working_dir, exist_ok=True)


class SuffixArraysPhraseTable(MosesFeature):
    injector_section = 'suffixarrays'
    injectable_fields = {
        'sample': ('number of samples for phrase table', int, 1000),
        'method': ('sampling method for suffix array phrase table', str, 'ranked3'),
    }

    def __init__(self, model, langs):
        MosesFeature.__init__(self, 'Mmsapt')

        self._sample = None  # Injected
        self._method = None  # Injected

        self._model = model
        self._source_lang = langs[0]
        self._target_lang = langs[1]

        self._symal_bin = os.path.join(Moses.bin_path, 'bin', 'symal')
        self._symal2mam_bin = os.path.join(Moses.bin_path, 'bin', 'symal2mam')
        self._mttbuild_bin = os.path.join(Moses.bin_path, 'bin', 'mtt-build')
        self._mmlexbuild_bin = os.path.join(Moses.bin_path, 'bin', 'mmlex-build')

    def _get_model_basename(self):
        return os.path.join(self._model, 'model')

    def get_iniline(self):
        template = 'path={model}. L1={source_lang} L2={target_lang} output-factor=0 sample={sample} method={method} ' \
                   'pfwd=g pbwd=g logcnt=0 coh=0 prov=0 rare=0 unal=0 smooth=.01 lexalpha=0 lr-func=DM0 bias-loglevel=1'

        return template.format(model=self.get_relpath(self._get_model_basename()),
                               source_lang=self._source_lang, target_lang=self._target_lang,
                               sample=self._sample, method=self._method)

    def train(self, corpora, aligner, working_dir='.', log_file=None):
        if os.path.isdir(self._model) and len(os.listdir(self._model)) > 0:
            raise Exception('Model already exists at ' + self._model)

        if not os.path.isdir(self._model):
            fileutils.makedirs(self._model, exist_ok=True)

        if not os.path.isdir(working_dir):
            fileutils.makedirs(working_dir, exist_ok=True)

        l1 = self._source_lang
        l2 = self._target_lang
        langs = (l1, l2)
        langs_suffix = l1 + '-' + l2

        mct_base = self._get_model_basename()
        dmp_file = mct_base + '.dmp'
        mam_file = mct_base + '.' + langs_suffix + '.mam'
        lex_file = mct_base + '.' + langs_suffix + '.lex'

        log = shell.DEVNULL

        try:
            if log_file is not None:
                log = open(log_file, 'a')

            merged_corpus = ParallelCorpus(os.path.basename(mct_base), working_dir, langs)

            # Create merged corpus and domains list file (dmp)
            fileutils.merge([corpus.get_file(l1) for corpus in corpora], merged_corpus.get_file(l1))
            fileutils.merge([corpus.get_file(l2) for corpus in corpora], merged_corpus.get_file(l2))
            with open(dmp_file, 'w') as dmp:
                for corpus in corpora:
                    dmp.write(str(corpus.name) + ' ' + str(corpus.count_lines()) + '\n')

            # Create alignments in 'bal' file and symmetrize
            bal_file = aligner.align(merged_corpus, langs, self._model, working_dir, log_file)

            symal_file = os.path.join(working_dir, 'alignments.' + langs_suffix + '.symal')
            symal_command = [self._symal_bin, '-a=g', '-d=yes', '-f=yes', '-b=yes']
            with open(bal_file) as stdin:
                with open(symal_file, 'w') as stdout:
                    shell.execute(symal_command, stdin=stdin, stdout=stdout, stderr=log)

            # Execute mtt-build
            mttbuild_command = self._get_mttbuild_command(mct_base, dmp_file, l1)
            with open(merged_corpus.get_file(l1)) as stdin:
                shell.execute(mttbuild_command, stdin=stdin, stdout=log, stderr=log)

            mttbuild_command = self._get_mttbuild_command(mct_base, dmp_file, l2)
            with open(merged_corpus.get_file(l2)) as stdin:
                shell.execute(mttbuild_command, stdin=stdin, stdout=log, stderr=log)

            # Create 'mam' file
            mam_command = [self._symal2mam_bin, mam_file]
            with open(symal_file) as stdin:
                shell.execute(mam_command, stdin=stdin, stdout=log, stderr=log)

            # Create 'lex' file
            lex_command = [self._mmlexbuild_bin, mct_base + '.', l1, l2, '-o', lex_file]
            shell.execute(lex_command, stdout=log, stderr=log)
        finally:
            if log_file is not None:
                log.close()

    def _get_mttbuild_command(self, mct_base, dmp_file, lang):
        output = mct_base + '.' + lang
        return [self._mttbuild_bin, '-i', '-o', output, '-m', dmp_file, '-g']


class FastAlign(WordAligner):
    def __init__(self):
        WordAligner.__init__(self)

        self._align_bin = os.path.join(scripts.BIN_DIR, 'fastalign-maurobuild', 'fast_align')

    def align(self, corpus, langs, model_dir, working_dir='.', log_file=None):
        WordAligner.align(self, corpus, langs, working_dir, log_file)

        l1 = langs[0]
        l2 = langs[1]
        corpus_name = 'corpus'
        langs_suffix = l1 + '-' + l2

        fwd_file = os.path.join(working_dir, corpus_name + '.' + langs_suffix + '.fwd')
        bwd_file = os.path.join(working_dir, corpus_name + '.' + langs_suffix + '.bwd')
        bal_file = os.path.join(working_dir, corpus_name + '.' + langs_suffix + '.bal')
        aligned_file_path = os.path.join(working_dir, corpus_name + '.' + langs_suffix + '.aligned')

        corpus_l1 = corpus.get_file(l1)
        corpus_l2 = corpus.get_file(l2)

        log = shell.DEVNULL

        try:
            if log_file is not None:
                log = open(log_file, 'a')

            with open(corpus_l1) as source_corpus, \
                    open(corpus_l2) as target_corpus, \
                    open(aligned_file_path, 'w') as aligned_file:
                for x, y in zip(source_corpus, target_corpus):
                    aligned_file.write(x.strip() + ' ||| ' + y.strip() + '\n')

            cpus = multiprocessing.cpu_count()

            # Forward alignments
            fwd_model = os.path.join(model_dir, 'model.align.fwd')
            command = [self._align_bin, '-d', '-v', '-o', '-n', str(cpus), '-B', '-p', fwd_model, '-i',
                       aligned_file_path]
            with open(fwd_file, 'w') as stdout:
                shell.execute(command, stdout=stdout, stderr=log)

            # Forward alignments
            bwd_model = os.path.join(model_dir, 'model.align.bwd')
            command = [self._align_bin, '-d', '-v', '-o', '-n', str(cpus), '-B', '-p', bwd_model, '-r', '-i',
                       aligned_file_path]
            with open(bwd_file, 'w') as stdout:
                shell.execute(command, stdout=stdout, stderr=log)
        finally:
            if log_file is not None:
                log.close()

        encoder = _FastAlignBALEncoder(corpus, langs, fwd_file, bwd_file)
        encoder.encode(bal_file)

        return bal_file


class _FastAlignBALEncoder:
    def __init__(self, corpus, langs, fwd, bwd):
        self._corpus_l1 = corpus.get_file(langs[0])
        self._corpus_l2 = corpus.get_file(langs[1])
        self._fwd = fwd
        self._bwd = bwd

    def encode(self, output):
        files = []

        try:
            files.append(open(self._corpus_l1))
            files.append(open(self._corpus_l2))
            files.append(open(self._fwd))
            files.append(open(self._bwd))
            files.append(open(output, 'w'))

            c1, c2, fwd, bwd, bal = files

            for t1 in c1:
                t1 = t1.strip().split()
{'fraction_non_alphanumeric': 0.12306943937735619, 'fraction_numerical': 0.007783047549556123, 'mean_word_length': 2.336105476673428, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1520510', 'n_tokens_mistral': 5655, 'n_tokens_neox': 5407, 'n_words': 1111}