Sentence Meaning and Natural Language Inference › Worked Example: A Small NLI Battery

import re

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

LABELS = ["contradiction", "entailment", "neutral"]
NEGATIONS = {
    "not",
    "no",
    "never",
    "none",
    "nobody",
    "nothing",
    "neither",
    "nor",
    "without",
}
ANTONYM_PAIRS = [
    ("open", "closed"),
    ("hot", "cold"),
    ("up", "down"),
    ("before", "after"),
    ("true", "false"),
    ("always", "never"),
    ("more", "less"),
    ("increase", "decrease"),
    ("buy", "sell"),
    ("win", "lose"),
]
STOPWORDS = {
    "a",
    "an",
    "the",
    "is",
    "are",
    "was",
    "were",
    "be",
    "been",
    "being",
    "of",
    "to",
    "in",
    "on",
    "at",
    "for",
    "with",
    "and",
    "or",
    "that",
    "this",
    "it",
    "as",
    "by",
    "from",
}


def simple_tokens(text):
    return re.findall(r"[a-z0-9]+", text.lower())

class LexicalNLIBaseline:
    """Fixed teaching heuristic, without training or model downloads."""

    def predict(self, pairs):
        rows = []
        for premise, hypothesis in pairs:
            pt = set(simple_tokens(premise))
            ht = set(simple_tokens(hypothesis))
            content = ht - STOPWORDS
            overlap = len(content & pt) / max(1, len(content))
            neg_p = any(token in NEGATIONS for token in pt)
            neg_h = any(token in NEGATIONS for token in ht)
            antonym_hit = any(
                (a in pt and b in ht) or (b in pt and a in ht)
                for a, b in ANTONYM_PAIRS
            )

            ent_logit = 0.2 + 1.6 * overlap
            con_logit = 0.1
            if neg_p != neg_h:
                con_logit += 1.1
            if antonym_hit:
                con_logit += 1.3
            neu_logit = 0.7 + 1.0 * (1.0 - overlap)

            logits = np.array([con_logit, ent_logit, neu_logit])
            probs = np.exp(logits - logits.max())
            rows.append(probs / probs.sum())
        return np.vstack(rows) if rows else np.empty((0, len(LABELS)))

nli_model = LexicalNLIBaseline()
model_labels = LABELS

def predict_nli_batch(premises, hypotheses):
    if len(premises) != len(hypotheses):
        raise ValueError("Each premise needs exactly one hypothesis.")
    pairs = list(zip(premises, hypotheses))
    raw = nli_model.predict(pairs)
    results = []
    for row in raw:
        probs = {label: float(row[i]) for i, label in enumerate(LABELS)}
        prediction = max(LABELS, key=lambda label: probs[label])
        results.append({"probs": probs, "prediction": prediction})
    return results


def predict_nli(premise, hypothesis):
    return predict_nli_batch([premise], [hypothesis])[0]

examples = [
    {
        "id": 1,
        "premise": "A young child is riding a bicycle in the park.",
        "hypothesis": "A child is outdoors.",
        "gold": "entailment",
    },
    {
        "id": 2,
        "premise": "The museum is open every day.",
        "hypothesis": "The museum is closed on Mondays.",
        "gold": "contradiction",
    },
    {
        "id": 3,
        "premise": "A woman is playing the violin.",
        "hypothesis": "A woman is performing in an orchestra.",
        "gold": "neutral",
    },
    {
        "id": 4,
        "premise": "All mammals are warm-blooded, and dolphins are mammals.",
        "hypothesis": "Dolphins are warm-blooded.",
        "gold": "entailment",
    },
    {
        "id": 5,
        "premise": "The train arrived before noon.",
        "hypothesis": "The train arrived after noon.",
        "gold": "contradiction",
    },
    {
        "id": 6,
        "premise": "The bank raised interest rates.",
        "hypothesis": "The river bank is steep.",
        "gold": "neutral",
    },
    {
        "id": 7,
        "premise": "No dogs are allowed in the store.",
        "hypothesis": "No puppies are allowed in the store.",
        "gold": "entailment",
    },
    {
        "id": 8,
        "premise": "The door is open.",
        "hypothesis": "The door is not open.",
        "gold": "contradiction",
    },
    {
        "id": 9,
        "premise": "The lawyer questioned the witness.",
        "hypothesis": "The witness questioned the lawyer.",
        "gold": "neutral",
    },
    {
        "id": 10,
        "premise": "Alice sold Bob a car.",
        "hypothesis": "Bob bought a car from Alice.",
        "gold": "entailment",
    },
    {
        "id": 11,
        "premise": "The keys are in the drawer.",
        "hypothesis": "The keys are in the kitchen.",
        "gold": "neutral",
    },
    {
        "id": 12,
        "premise": "The meeting starts at 3 p.m.",
        "hypothesis": "The meeting starts at 4 p.m.",
        "gold": "contradiction",
    },
]
examples_df = pd.DataFrame(examples)

batch_predictions = predict_nli_batch(
    examples_df["premise"].tolist(),
    examples_df["hypothesis"].tolist(),
)

results_df = examples_df.copy()
results_df["prediction"] = [item["prediction"] for item in batch_predictions]
for label in LABELS:
    results_df[f"p_{label}"] = [
        item["probs"][label] for item in batch_predictions
    ]
results_df["correct"] = results_df["gold"] == results_df["prediction"]

directional_pairs = [
    ("A poodle barked.", "A dog barked."),
    ("A dog barked.", "A poodle barked."),
    ("The restaurant is open.", "The restaurant is not closed."),
    ("The restaurant is not closed.", "The restaurant is open."),
]

directional_preds = predict_nli_batch(
    [p for p, h in directional_pairs],
    [h for p, h in directional_pairs],
)

directional_df = pd.DataFrame(
    directional_pairs, columns=["premise", "hypothesis"]
)
directional_df["prediction"] = [
    item["prediction"] for item in directional_preds
]
directional_df["p_entailment"] = [
    item["probs"]["entailment"] for item in directional_preds
]
directional_df["p_neutral"] = [
    item["probs"]["neutral"] for item in directional_preds
]
directional_df["p_contradiction"] = [
    item["probs"]["contradiction"] for item in directional_preds
]

presupposition_pairs = [
    ("The king of France is bald.", "There is a king of France."),
    ("The king of France is not bald.", "There is a king of France."),
    ("John stopped smoking.", "John used to smoke."),
    ("John did not stop smoking.", "John used to smoke."),
]

presupposition_preds = predict_nli_batch(
    [p for p, h in presupposition_pairs],
    [h for p, h in presupposition_pairs],
)

presupposition_df = pd.DataFrame(
    presupposition_pairs, columns=["premise", "hypothesis"]
)
presupposition_df["prediction"] = [
    item["prediction"] for item in presupposition_preds
]
presupposition_df["p_entailment"] = [
    item["probs"]["entailment"] for item in presupposition_preds
]
presupposition_df["p_neutral"] = [
    item["probs"]["neutral"] for item in presupposition_preds
]
presupposition_df["p_contradiction"] = [
    item["probs"]["contradiction"] for item in presupposition_preds
]

stress_examples = [
    {
        "premise": "The lawyer questioned the witness.",
        "hypothesis": "The witness questioned the lawyer.",
        "gold": "neutral",
    },
    {
        "premise": "The lawyer questioned the witness.",
        "hypothesis": "The lawyer questioned someone.",
        "gold": "entailment",
    },
    {
        "premise": "The door is open.",
        "hypothesis": "The door is closed.",
        "gold": "contradiction",
    },
    {
        "premise": "The door is not open.",
        "hypothesis": "The door is closed.",
        "gold": "entailment",
    },
    {
        "premise": "No dogs are allowed in the store.",
        "hypothesis": "No puppies are allowed in the store.",
        "gold": "entailment",
    },
    {
        "premise": "No puppies are allowed in the store.",
        "hypothesis": "No dogs are allowed in the store.",
        "gold": "neutral",
    },
]

stress_df = pd.DataFrame(stress_examples)
stress_preds = predict_nli_batch(
    stress_df["premise"].tolist(),
    stress_df["hypothesis"].tolist(),
)
stress_df["prediction"] = [item["prediction"] for item in stress_preds]
stress_df["correct"] = stress_df["gold"] == stress_df["prediction"]

confusion = np.zeros((len(LABELS), len(LABELS)), dtype=int)
for gold, pred in zip(results_df["gold"], results_df["prediction"]):
    confusion[LABELS.index(gold), LABELS.index(pred)] += 1


def macro_f1(y_true, y_pred, labels):
    scores = []
    for label in labels:
        tp = sum(t == label and p == label for t, p in zip(y_true, y_pred))
        fp = sum(t != label and p == label for t, p in zip(y_true, y_pred))
        fn = sum(t == label and p != label for t, p in zip(y_true, y_pred))
        precision = tp / (tp + fp) if tp + fp else 0.0
        recall = tp / (tp + fn) if tp + fn else 0.0
        f1 = (
            2 * precision * recall / (precision + recall)
            if precision + recall
            else 0.0
        )
        scores.append(f1)
    return float(np.mean(scores))


accuracy = float((results_df["gold"] == results_df["prediction"]).mean())
f1 = macro_f1(
    results_df["gold"].tolist(), results_df["prediction"].tolist(), LABELS
)

plot_ids = [1, 2, 4, 7]
plot_df = results_df[results_df["id"].isin(plot_ids)].copy().sort_values("id")