def with_lookup(old_fn, lookup):
    """
    Wrap a similarity function, looking up a pred for each given lemma,
    returning 0 if the lemma was not found
    :param old_fn: function taking two indices and returning a similarity score
    :param lookup: mapping from lemmas to sets of indices
    :return: similarity function
    """
    def sim_fn(a, b, *args, **kwargs):
        "Calculate similarity"
        try:
            return old_fn(lookup[a], lookup[b], *args, **kwargs)
        except KeyError:
            # Unknown lemmas
            return 0
    return sim_fn