import math


def cooccur_probability_fn(counts):
  """Returns function for computing entailment probability.

  Args:
    counts: Dict mapping unigrams / bigrams (joined using "|||") to their
      counts.

  Returns:
    Function handle to compute entailment probability.
  """

  def _cooccur_probability(ngram, table):
    """Returns probability of ngram being entailed by the table.

    Uses the co-occurrence counts given along with the lexical
    entailment model described in:

      Glickman, Oren, Ido Dagan, and Moshe Koppel.
      "A lexical alignment model for probabilistic textual entailment."
      Machine Learning Challenges.
      Springer, Berlin, Heidelberg, 2006. 287-298.

    E.g.:
      >>> _cooccur_probability(["michael", "dahlquist"],
                                  [(["name"], ["michael", "dahlquist"])])
      >>> 1.0

    Args:
      ngram: List of tokens.
      table: List of either (attribute, value) pairs or (head, relation, tail)
        triples. Each member of the pair / triple is assumed to already be
        tokenized into a list of strings.

    Returns:
      prob: Float probability of ngram being entailed by the table.
    """
    table_toks = set()
    for item in table:
      if len(item) == 2:
        # attribute, value
        table_toks.add("_".join(item[0]))
        table_toks.update(item[1])
      else:
        # head, relation, tail
        table_toks.update(item[0] + ["_".join(item[1])] + item[2])
    probability = 1.
    for xtok in ngram:
      if xtok in table_toks:
        continue
      max_p = 0.
      for btok in table_toks:
        if btok not in counts:
          continue
        p = float(counts.get(btok + "|||" + xtok, 0.)) / counts[btok]
        if p > max_p:
          max_p = p
      probability *= max_p
    return math.pow(probability, 1. / len(ngram))

  return _cooccur_probability