def smart_ropen(path):
    """Opens files directly or through gzip depending on extension."""
    if path.endswith('.gz'):
        return TextIOWrapper(gzip.open(path, 'rb'))
    else:
        return open(path, 'r') def smart_wopen(path):
    """Opens files directly or through gzip depending on extension."""
    if path.endswith('.gz'):
        return TextIOWrapper(gzip.open(path, 'wb'))
    else:
        return open(path, 'w') def make_unique_directory(dir=None):
    return tempfile.mkdtemp(prefix=datetime.now().strftime("%y%m%d_%H%M%S_"), dir=dir) def pickle_it(path, obj):
    """Dump a pickled representation of the object to disk"""
    with open(path, 'wb') as fo:
        pickle.dump(obj, fo) def unpickle_it(path):
    """Load an object from a pickled representation stored on disk"""
    with open(path, 'rb') as fo:
        return pickle.load(fo) def traceit(func):
    @wraps(func)
    def newfunc(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except:
            raise Exception(''.join(traceback.format_exception(*sys.exc_info())))
    return newfunc def make_models(wmap, extractors: 'list[Extractor]', uniform_weights=False) -> ModelContainer:
    wmap = dict(wmap)
    # all scorers sorted by id
    extractors = tuple(sorted(extractors, key=lambda x: x.id))

    if uniform_weights:
        for extractor in extractors:
            fnames = extractor.fnames(wmap.keys())
            for fname in fnames:
                wmap[fname] = 1.0 / len(extractors) / len(fnames)

    return ModelContainer(wmap, extractors) def read_weights(path, default=None, random=False, temperature=1.0, u=0, std=0.01):
    """
    Read a sequence of key-value pairs.
    :param path: file where to read sequence from
    :param default: if set, overwrites the values read from file
    :param random: if set, sample values from N(u, std)
    :param temperature: scales the final weight: weight/T
    :param u: mean of normal
    :param std: standard deviation
    :return:
    """
    wmap = {}
    with smart_ropen(path) as fi:
        for line in fi.readlines():
            fields = line.split()
            if len(fields) != 2:
                continue
            w = float(fields[1])
            if default is not None:
                w = default
            elif random:
                w = np.random.normal(u, std)
            w /= temperature
            wmap[fields[0]] = w
    return wmap def load_grammar(path):
    """
    Load a grammar from a text file.
    :param path:
    :return:
    """
    return SCFG(iterrules(smart_ropen(path))) class GoalRuleMaker:

    def __init__(self, goal_str, start_str, n=0):
        self._n = n
        self._goal_str = goal_str
        self._start_str = start_str

    def _make_goal_str(self, n=None):
        return '{0}{1}'.format(self._goal_str, self._n) if n is None else '{0}{1}'.format(self._goal_str, n)

    def update(self):
        """Updates the state of the factory"""
        self._n += 1

    def get_srule(self):
        """Returns that goal rule based on the state of the factory"""
        if self._n == 0:
            rhs = [Nonterminal(self._start_str)]
        else:
            rhs = [Nonterminal(self._make_goal_str(self._n - 1))]
        return SCFGProduction(Nonterminal(self._make_goal_str()), rhs, rhs, [1], {'GoalRule': 1.0})

    def get_next_srule(self):
        """Returns what would be the next goal rule without updating the state of the factory."""
        rhs = [Nonterminal(self._make_goal_str(self._n))]
        return SCFGProduction(Nonterminal(self._make_goal_str(self._n + 1)), rhs, rhs, [1], {'GoalRule': 1.0})

    def get_iview(self):
        return InputGroupView([self.get_srule()])

    def get_oview(self):
        return OutputView(self.get_srule()) def save_forest(path, hg, ffs=None):
    with smart_wopen(path) as fo:
        print('# FOREST', file=fo)
        print(hg, file=fo)
        if ffs is not None:
            print('# FF')
            for ff in ffs:
                print(ff, file=fo) def save_ffs(path, ffs):
    with open(path, 'wb') as fo:
        pickle.dump(ffs, fo) def load_ffs(path):
    with open(path, 'rb') as fo:
        return pickle.load(fo) def make_dead_srule(lhs='X', dead='<dead-end>', fname='DeadRule'):
    return SCFGProduction(Nonterminal(lhs), (Terminal(dead),), (Terminal(dead),), [], {fname: 1.0}) def make_batches(data, sizeperc):
    """

    :param data:
    :param sizeperc: size of a batche expressed as a percentage of the total data
        (between 0 and 1)
    :return:
    """
    if sizeperc > 1:
        raise ValueError('sizeperc must be between 0 and 1')
    if sizeperc < 0:
        raise ValueError('sizeperc must be between 0 and 1')
    b_size = max(int(len(data) * sizeperc), 1)
    n_batches = max(int(len(data)/b_size), 1)
    batches = [[] for _ in range(n_batches)]
    pool = deque(data)
    b = 0
    while pool:
        datum = pool.popleft()
        batches[b].append(datum)
        b += 1
        if b == n_batches:
            b = 0
    return batches def number_of_batches(total, sizeperc):
    """

    :param data:
    :param sizeperc: size of a batche expressed as a percentage of the total data
        (between 0 and 1)
    :return:
    """
    if sizeperc > 1:
        raise ValueError('sizeperc must be between 0 and 1')
    if sizeperc < 0:
        raise ValueError('sizeperc must be between 0 and 1')
    b_size = max(int(total * sizeperc), 1)
    return max(int(total/b_size), 1) class SegmentMetaData(object):
    """
    A simple container for input segments
    """

    def __init__(self, sid, src, grammar, refs=[]):
        self.sid_ = sid
        self.grammar_ = grammar
        self.src_ = src
        self.refs_ = tuple(refs)

    @property
    def id(self):
        return self.sid_

    @property
    def src(self):
        return self.src_

    def src_tokens(self, boundaries=False):
        if not boundaries:
            return self.src.split()
        else:
            return ['<s>'] + self.src.split() + ['</s>']

    @property
    def refs(self):
        return self.refs_

    @property
    def grammar(self):
        return self.grammar_

    def __str__(self):
        return 'grammar=%s\tsrc=%s' % (self.grammar_, self.src_)

    def to_sgm(self, dump_refs=True):
        if dump_refs and self.refs_:
            return '<seg grammar="{1}" id="{0}">{2}</seg> ||| {3}'.format(self.sid_,
                                                                          self.grammar_,
                                                                          self.src_,
                                                                          ' ||| '.join(str(ref) for ref in self.refs_))
        else:
            return '<seg grammar="{1}" id="{0}">{2}</seg>'.format(self.sid_,
                                                                  self.grammar_,
                                                                  self.src_)

    @staticmethod
    def parse(line, sid=None, grammar_dir=None, mode='cdec-sgml'):
        if mode == 'cdec-sgml':
            args = parse_cdec_sgml(line)
        else:
            raise Exception('unknown input format: %s' % mode)
        # overrides sentence id
        if sid is not None:
            args['sid'] = sid
        else:
            sid = args['sid']
        # overrides grammar
        if grammar_dir is not None:
            args['grammar'] = '{0}/grammar.{1}.gz'.format(grammar_dir, args['sid'])
        else:
            grammar_dir = args['grammar']
        # sanity checks
        if not isfile(args['grammar']):
            # a few more chances
            failed = False
            path = args['grammar']
            if path.endswith('.gz'):
                path = path[:-3]
                if isfile(path):
                    args['grammar'] = path
                else:
                    failed = True
            else:
                path += '.gz'
                if isfile(path):
                    args['grammar'] = path
                else:
                    failed = True
            if failed:
                raise FileNotFoundError('Grammar file not found: %s' % args['grammar'])
        # construct segment
        return SegmentMetaData(**args) def make_pass_grammar(seg, grammars, semiring, unk_lhs):
    """
    Make an input fsa for an input segment as well as its pass-through grammar.
    :param seg: a Segment object.
    :param grammars: a sequence of SCFGs.
    :param semiring: must provide `one`.
    :return: the input WDFSA, the pass-through grammar
    """
    fsa = WDFSA()
    pass_grammar = SCFG()
    unk = Nonterminal(unk_lhs)
    tokens = seg.src.split()
    for i, token in enumerate(tokens):
        word = Terminal(token)
        if any(g.in_ivocab(word) for g in grammars):
            pass_grammar.add(SCFGProduction.create(unk,
                                                   [word],
                                                   [word],
                                                   {'PassThrough': 1.0}))
        else:
            pass_grammar.add(SCFGProduction.create(unk,
                                                   [word],
                                                   [word],
                                                   {'PassThrough': 1.0,
                                                    'Unknown': 1.0}))
        fsa.add_arc(i, i + 1, word, semiring.one)
    fsa.make_initial(0)
    fsa.make_final(len(tokens))
    return fsa, pass_grammar def save_mcmc_yields(path, groups):
    """

    :param path: where to save
    :param groups: sorted list of sampled (obtained by group_by_projection)
    """
    with smart_wopen(path) as out:
        total = sum(group.count for group in groups)
        print('# MCMC samples={0}\n# estimate\tcount\tderivations\tyield'.format(total), file=out)
        for i, group in enumerate(groups, 1):
            print('{0}\t{1}\t{2}\t{3}'.format(float(group.count)/total,
                                              group.count,
                                              len(set(group.values)),
                                              group.key),
                  file=out) def save_mcmc_derivations(path, groups, valuefunc, derivation2str, compfunc=None):  # =lambda d: DerivationYield.derivation(d.rules())
    """

    :param path: where to save
    :param samples: sorted list of samples (obtained by group_by_identity)
    :param valuefunc: compute the value of a derivation
    :param omega_d: a function over derivations
    """
    with smart_wopen(path) as out:
        total = sum(group.count for group in groups)
        print('# MCMC samples={0}'.format(total), file=out)
        if compfunc is not None:
            print('# estimate\tcount\tscore\tderivation\tfeatures', file=out)
            for i, group in enumerate(groups, 1):
                sample = group.key
                print('{0}\t{1}\t{2}\t{3}\t{4}'.format(group.count / total,  # estimate
                                                       group.count,
                                                       valuefunc(sample),
                                                       derivation2str(sample),
                                                       compfunc(sample)),
                      file=out)
        else:
            print('# estimate\tcount\tscore\tderivation', file=out)
            for i, group in enumerate(groups, 1):
                sample = group.key
                print('{0}\t{1}\t{2}\t{3}'.format(group.count/total,  # estimate
                                                  group.count,
                                                  valuefunc(sample),
                                                  derivation2str(sample)),
                      file=out) def save_markov_chain(path, markov_chain, derivation2str, valuefunc=None, compfunc=None, flat=True):
    """

    :param path: where to save
    :param markov_chain: the original Markov chain
    :param valuefunc: an optional function over derivations that returns a score
    :param compfunc: an optional function over derivations that return feature components
    :param flat: whether the Markov chain is flat (each state represents a single derivation) or not,
        in which case each state is a sequence of derivations.
    """
    if flat:
        with smart_wopen(path) as out:
            for d in markov_chain:
                fields = []
                if valuefunc is not None:
                    fields.append(valuefunc(d))
                fields.append(derivation2str(d))
                if compfunc is not None:
                    fields.append(compfunc(d))
                print('\t'.join(str(x) for x in fields), file=out)
    else:
        with smart_wopen(path) as out:
            for i, batch in enumerate(markov_chain):
                for d in batch:
                    fields = [i]
                    if valuefunc is not None:
                        fields.append(valuefunc(d))
                    fields.append(derivation2str(d))
                    if compfunc is not None:
                        fields.append(compfunc(d))
                    print('\t'.join(str(x) for x in fields), file=out) def read_segments_from_file(path, grammar_dir=None, shuffle=False) -> 'tuple':
    """
    Read cdec-formated input segments (possibly along with their reference translations) from a file.
    :param path: path to file (possibly gzipped)
    :param grammar_dir: overwrites grammar directory
    :param shuffle: shuffle segments inplace
    :return: tuple of SegmentMetaData objects
    """
    return read_segments_from_stream(smart_ropen(path), grammar_dir=grammar_dir, shuffle=shuffle) def save_segments(path, segments):
    with smart_wopen(path) as fo:
        for seg in segments:
            print(seg.to_sgm(True), file=fo) def load_feature_extractors(rt=None, wp=None, ap=None, slm=None, lm=None) -> 'tuple':
    """
    Load feature extractors depending on command line options.

    For now we have the following extractors:

        * RuleTable
        * WordPenalty
        * ArityPenalty
        * StatelessLM
        * KenLM

    :return: a tuple of Extractor objects
    """
    extractors = []

    if rt:
        extractor = RuleTable(uid=len(extractors),
                              name='RuleTable')
        extractors.append(extractor)

    if wp:
        extractor = WordPenalty(uid=len(extractors),
                                name=wp[0],
                                penalty=float(wp[1]))
        extractors.append(extractor)

    if ap:
        extractor = ArityPenalty(uid=len(extractors),
                                 name=ap[0],
                                 penalty=float(ap[1]))
        extractors.append(extractor)

    if slm:
        extractor = StatelessLM(uid=len(extractors),
                                name=slm[0],
                                order=int(slm[1]),
                                path=slm[2])
        extractors.append(extractor)

    if lm:
        extractor = KenLM(uid=len(extractors),
                          name=lm[0],
                          order=int(lm[1]),
                          path=lm[2])
        extractors.append(extractor)

    return tuple(extractors) @traceit
def training_decode(seg, args, n_samples, staticdir, decisiondir, model, redo, log=dummyfunc):

    saving = {
        'forest': '{0}/{1}.hyp.forest'.format(staticdir, seg.id),
        'lookup': '{0}/{1}.hyp.ffs.rule'.format(staticdir, seg.id),
        'stateless': '{0}/{1}.hyp.ffs.stateless'.format(staticdir, seg.id),
        'decisions': '{0}/{1}.gz'.format(decisiondir, seg.id)
    }
    return decode(seg, args, n_samples, model, saving, redo, log)
import logging
import argparse
import sys
import subprocess as sp
import shlex
import argparse
import logging
import sys
import itertools
import os
import numpy as np
import traceback
import grasp.ptypes as ptypes
import grasp.semiring as semiring
from os.path import splitext
from multiprocessing import Pool
from functools import partial
from collections import deque
from grasp.loss.fast_bleu import DecodingBLEU
from grasp.recipes import smart_ropen, smart_wopen, make_unique_directory, pickle_it, unpickle_it, traceit
from grasp.scoring.scorer import TableLookupScorer, StatelessScorer, StatefulScorer
from grasp.scoring.util import make_models
from grasp.scoring.util import read_weights
from grasp.mt.cdec_format import load_grammar
from grasp.mt.util import GoalRuleMaker
from grasp.mt.util import save_forest, save_ffs, load_ffs, make_dead_srule, make_batches, number_of_batches
from grasp.mt.segment import SegmentMetaData
from grasp.mt.input import make_pass_grammar
from grasp.semiring.operator import FixedLHS, FixedRHS
from grasp.formal.scfgop import output_projection
from grasp.formal.fsa import make_dfa, make_dfa_set, make_dfa_set2
from grasp.formal.scfgop import make_hypergraph_from_input_view, output_projection
from grasp.formal.scfgop import lookup_components, stateless_components
from grasp.formal.topsort import AcyclicTopSortTable
from grasp.formal.traversal import bracketed_string, yield_string
from grasp.formal.wfunc import TableLookupFunction, ConstantFunction, derivation_weight
from grasp.cfg.model import DummyConstant
from grasp.cfg.symbol import Nonterminal
from grasp.cfg.symbol import Terminal
from grasp.cfg.srule import OutputView
from grasp.alg.deduction import NederhofParser, EarleyParser, EarleyRescorer
from grasp.alg.inference import viterbi_derivation, AncestralSampler
from grasp.alg.value import acyclic_value_recursion, acyclic_reversed_value_recursion, compute_edge_expectation
from grasp.alg.rescoring import weight_edges
from grasp.alg.rescoring import SlicedRescoring
from grasp.alg.rescoring import stateless_rescoring
from grasp.alg.chain import apply_filters, group_by_identity, group_by_projection
from grasp.alg.expectation import expected_components
from grasp.scoring.frepr import FComponents
from grasp.io.results import save_mcmc_yields, save_mcmc_derivations, save_markov_chain
from random import shuffle
from numpy import linalg as LA
from scipy.optimize import minimize
from time import time, strftime
from types import SimpleNamespace
from grasp.mt.pipeline import read_segments_from_file
from grasp.mt.pipeline import save_segments
from grasp.mt.pipeline import load_feature_extractors
from grasp.mt.pipeline import training_decode
    from grasp.mt.pipeline import training_biparse

        N is the total number of training instances (not only those in this batch)
        M is the batch size
        (x_i: for 1 <= i <= M) are the training instances in the batch
        \lambda is a L2 regularisation strength
        \gamma is a learning rate


    :param args:
    :param staticdir:
    :param workspace:
    :param model:
    :param iteration:
    :param batch_number:
    :param batch:
    :param merging:
    :param training_size:
    :return:
    """
    logging.info('[I=%d] batch=%d size=%d total=%d gamma=%f', iteration, batch_number, len(batch), N, gamma)

    def f(theta):

        model_t = make_models(dict(zip(model.fnames(), theta)), model.extractors())

        logprob, derivatives = objective_and_derivatives(args, staticdir, workspace, model_t, batch, merging)
        obj = logprob
        jac = np.array(list(derivatives.densify()), dtype=ptypes.weight)

        if args.L2 == 0.0:
            logging.info('[I=%d, b=%d] O=%f', iteration, batch_number, obj)
            logging.debug('[I=%d, b=%d] Derivatives\n%s', iteration, batch_number, npvec2str(jac, model.fnames(), '\n'))
            return obj, jac
        else:
            r_obj = obj
            r_jac = jac.copy()

            if args.L2 != 0.0:  # L2-regularised:  p(theta_k) \propto exp(-\lambda_k ^ 2 / (2*\sigma^2), sigma^2=1, lambda_k = lambda
                regulariser = LA.norm(theta, 2) ** 2
                r_obj -= args.L2 * regulariser / (2 * N)  #  0.5 * lambda * ||theta||^2
                r_jac -= args.L2 * theta / N  # lambda * theta
                logging.info('[I=%d, b=%d] O=%f L=%f', iteration, batch_number, r_obj, obj)
                logging.debug('[I=%d, b=%d] Derivatives\n%s', iteration, batch_number, npvec2str(r_jac, model.fnames()))

            return r_obj, r_jac

    w0 = np.array(list(model.weights().densify()), dtype=ptypes.weight)
    obj_b, jac_b = f(w0)
    w1 = w0 - gamma * jac_b
    # TODO: AdaGrad
    # TODO: Finite differences

    return w1


def eval_devtest(args, workspace, iterdir, model, segments):
    logging.info('Assessing loss on validation set')
    evaldir = '{0}/mteval'.format(iterdir)
    os.makedirs(evaldir, exist_ok=True)
    staticdir = '{0}/{1}'.format(workspace, args.devtest_alias)
    bleu = mteval(args, staticdir, model, segments,
                  '{0}/{1}.hyps'.format(evaldir, args.devtest_alias),
                  '{0}/refs'.format(staticdir, args.devtest_alias),
                  '{0}/{1}.bleu'.format(evaldir, args.devtest_alias),
                  '{0}/{1}.decisions'.format(evaldir, args.devtest_alias))
    logging.info('BLEU %s %s', args.devtest_alias, bleu)
    return bleu


def parse_training(args, staticdir, model, segments):
    logging.info('Parsing %d training instances using %d workers', len(segments), args.jobs)


    with Pool(args.jobs) as workers:
        feedback = workers.map(partial(training_biparse,
                                       args=args,
                                       workingdir=staticdir,
                                       model=model,
                                       log=logging.info),
                               segments)
    return tuple([seg for seg, status in zip(segments, feedback) if status])


def sanity_checks(args):
    failed = False
    if not os.path.exists(args.scoring_tool):
        logging.error('Scoring tool not found: %s', args.scoring_tool)
        failed = True
    if not os.path.exists(args.dev):
        logging.error('Training set not found: %s', args.dev)
        failed = True
    if args.devtest and not os.path.exists(args.devtest):
        logging.error('Validation set not found: %s', args.devtest)
        failed = True
    if not os.path.exists(args.weights):
        logging.error('Model description not found: %s', args.weights)
        failed = True
    return not failed


def core(args):

    workspace, devdir = make_dirs(args)

    if not sanity_checks(args):
        raise FileNotFoundError('One or more files could not be found')

    # 1. Make model
    logging.info('Loading feature extractors')
    # Load feature extractors
    extractors = load_feature_extractors(rt=args.rt, wp=args.wp, ap=args.ap, slm=args.slm, lm=args.lm)
    logging.info('Making model')
    # Load the model
    if args.init is None:
        model = make_models(read_weights(args.weights, temperature=args.temperature),
                            extractors)
    elif args.init == 'random':
        model = make_models(read_weights(args.weights, random=True, temperature=args.temperature),
                            extractors)
    elif args.init == 'uniform':

        model = make_models(read_weights(args.weights, temperature=args.temperature),