cpllab/syntaxgym

1from typing import Union, Optional as TOptional, List as TList2 3 4from pyparsing import *5import numpy as np6 7METRICS = {8    'sum': sum,9    'mean': np.mean,10    'median': np.median,11    'range': np.ptp,12    'max': max,13    'min': min14}15 16 17# Enable parser packrat (caching)18ParserElement.enablePackrat()19 20# Relative and absolute tolerance thresholds for surprisal equality21EQUALITY_RTOL = 1e-522EQUALITY_ATOL = 1e-323 24 25#######26# Define a grammar for prediction formulae.27 28# References a surprisal region29lpar = Suppress("(")30rpar = Suppress(")")31region = lpar + (Word(nums) | "*") + Suppress(";%") + Word(alphanums + "_-") + Suppress("%") + rpar32literal_float = pyparsing_common.number33 34class Region(object):35    def __init__(self, tokens):36        self.region_number = tokens[0]37        self.condition_name = tokens[1]38 39    def __str__(self):40        return "(%s;%%%s%%)" % (self.region_number, self.condition_name)41 42    def __repr__(self):43        return "Region(%s,%s)" % (self.condition_name, self.region_number)44 45    def __call__(self, surprisal_dict):46        if self.region_number == "*":47            return sum(value for (condition, region), value in surprisal_dict.items()48                       if condition == self.condition_name)49 50        return surprisal_dict[self.condition_name, int(self.region_number)]51 52class LiteralFloat(object):53    def __init__(self, tokens):54        self.value = float(tokens[0])55 56    def __str__(self):57        return "%f" % (self.value,)58 59    def __repr__(self):60        return "LiteralFloat(%f)" % (self.value,)61 62    def __call__(self, surprisal_dict):63        return self.value64 65class BinaryOp(object):66    operators: TOptional[TList[str]]67 68    def __init__(self, tokens):69        self.operator = tokens[0][1]70        if self.operators is not None and self.operator not in self.operators:71            raise ValueError("Invalid %s operator %s" % (self.__class__.__name__,72                                                            self.operator))73        self.operands = [tokens[0][0], tokens[0][2]]74 75    def __str__(self):76        return "(%s %s %s)" % (self.operands[0], self.operator, self.operands[1])77 78    def __repr__(self):79        return "%s(%s)(%s)" % (self.__class__.__name__, self.operator, ",".join(map(repr, self.operands)))80 81    def __call__(self, surprisal_dict):82        op_vals = [op(surprisal_dict) for op in self.operands]83        return self._evaluate(op_vals, surprisal_dict)84 85    def _evaluate(self, evaluated_operands, surprisal_dict):86        raise NotImplementedError()87 88class BoolOp(BinaryOp):89    operators = ["&", "|"]90    def _evaluate(self, op_vals, surprisal_dict):91        if self.operator == "&":92            return op_vals[0] and op_vals[1]93        elif self.operator == "|":94            return op_vals[0] or op_vals[1]95 96class FloatOp(BinaryOp):97    operators = ["-", "+"]98    def _evaluate(self, op_vals, surprisal_dict):99        if self.operator == "-":100            return op_vals[0] - op_vals[1]101        elif self.operator == "+":102            return op_vals[0] + op_vals[1]103 104class ComparatorOp(BinaryOp):105    operators = ["<", ">", "="]106    def _evaluate(self, op_vals, surprisal_dict):107        if self.operator == "<":108            return op_vals[0] < op_vals[1]109        elif self.operator == ">":110            return op_vals[0] > op_vals[1]111        elif self.operator == "=":112            return np.isclose(op_vals[0], op_vals[1],113                                rtol=EQUALITY_RTOL,114                                atol=EQUALITY_ATOL)115 116def Chain(op_cls, left_assoc=True):117    def chainer(tokens):118        """119        Create a binary tree of BinaryOps from the given repeated application120        of the op.121        """122        operators = tokens[0][1::2]123        args = tokens[0][0::2]124        if not left_assoc:125            raise NotImplementedError126 127        arg1 = args.pop(0)128        while len(args) > 0:129            operator = operators.pop(0)130            arg2 = args.pop(0)131            arg1 = op_cls([[arg1, operator, arg2]])132 133        return arg1134 135    return chainer136 137atom = region.setParseAction(Region) | literal_float.setParseAction(LiteralFloat)138 139prediction_expr = infixNotation(140    atom,141    [142        (oneOf("- +"), 2, opAssoc.LEFT, Chain(FloatOp)),143        (oneOf("< > ="), 2, opAssoc.LEFT, ComparatorOp),144        (oneOf("& |"), 2, opAssoc.LEFT, Chain(BoolOp)),145    ],146    lpar=lpar, rpar=rpar147)148 149 150class Prediction(object):151    """152    Predictions state expected relations between language model surprisal153    measures in different regions and conditions of a test suite. For more154    information, see :ref:`architecture`.155    """156 157    def __init__(self, idx: int, formula: Union[str, BinaryOp], metric: str):158        """159        Args:160            idx: A unique prediction ID. This is only relevant for161                serialization.162            formula: A string representation of the prediction formula, or an163                already parsed formula. For more information, see164                :ref:`architecture`.165            metric: Metric for aggregating surprisals within regions.166        """167        if isinstance(formula, str):168            try:169                formula = prediction_expr.parseString(formula, parseAll=True)[0]170            except ParseException as e:171                raise ValueError("Invalid formula expression %r" % (formula,)) from e172 173        self.idx = idx174        self.formula = formula175 176        if metric not in METRICS.keys():177            raise ValueError("Unknown metric %s. Supported metrics: %s" %178                             (metric, " ".join(METRICS.keys())))179        self.metric = metric180 181    def __call__(self, item):182        """183        Evaluate the prediction on the given item dict representation. For more184        information on item representations, see :ref:`suite_json`.185        """186        # Prepare relevant surprisal dict187        surps = {(c["condition_name"], r["region_number"]): r["metric_value"][self.metric]188                 for c in item["conditions"]189                 for r in c["regions"]}190        return self.formula(surps)191 192    @classmethod193    def from_dict(cls, pred_dict, idx: int, metric: str):194        """195        Parse from a prediction dictionary representation (see196        :ref:`suite_json`).197        """198        if not pred_dict["type"] == "formula":199            raise ValueError("Unknown prediction type %s" % (pred_dict["type"],))200 201        return cls(formula=pred_dict["formula"], idx=idx, metric=metric)202 203    @property204    def referenced_regions(self):205        """206        Get a set of the regions referenced by this formula.207        Each item is a tuple of the form ``(condition_name, region_number)``.208        """209        def traverse(x, acc):210            if isinstance(x, BinaryOp):211                for val in x.operands:212                    traverse(val, acc)213            elif isinstance(x, Region):214                acc.add((x.condition_name, int(x.region_number)))215 216            return acc217 218        return traverse(self.formula, set())219 220    def as_dict(self):221        """222        Serialize as a prediction dictionary representation (see223        :ref:`suite_json`).224        """225        return dict(type="formula", formula=str(self.formula))226 227    def __str__(self):228        return "Prediction(%s)" % (self.formula,)229    __repr__ = __str__230 231    def __hash__(self):232        return hash(self.formula)233 234    def __eq__(self, other):235        return isinstance(other, Prediction) and hash(self) == hash(other)236