Can DSPy be used with a discriminative NLI model like BERT

[umarbutler]

Suppose we have a BERT NLI model trained to zero-shot classify texts given a prompt such as, for example, 'This text relates to cars.' Is it possible to use DSPy to optimise the prompt for such a model?

[okhat]

Thanks @umarbutler ! I initially misunderstood the question.

DSPy doesn't currently support BERT-like models (any non-autoregressive models) natively.

However, it's indeed possible and semi-straightforward to set up prompt optimization for BERT-style models. It's just going to require one level of indirection.

import dspy
from typing import Callable, Literal

NLI = Literal["entailment", "contradiction", "neutral"]

class PrompterNLI(dspy.Signature):
    premise: str = dspy.InputField()
    hypothesis: str = dspy.InputField()
    prompt: str = dspy.OutputField(desc="prompt for a BERT model that conducts NLI")


class PromptedBERT(dspy.Module):
    def __init__(self, bert_nli_classifier: Callable[[str], NLI]):
        self.bert_nli_classifier = bert_nli_classifier
        self.prompter = dspy.ChainOfThought(PrompterNLI)
    
    def forward(self, premise: str, hypothesis: str) -> NLI:
        prompt = self.prompter(premise=premise, hypothesis=hypothesis)
        return self.bert_nli_classifier(prompt)

optimizer = dspy.MIPROv2(metric=METRIC, auto="medium", num_threads=THREADS)
optimized_bert = optimizer.compile(PromptedBERT(), trainset=TRAINSET, requires_permission_to_run=False)

Note that you need to specify METRIC, THREADS, TRAINSET, and build a nice wrapper bert_nli_classifier: Callable[[str], NLI] to use this.

Now here's the catch: the self.prompter currently depends on the premise and hypotheses, which means it'll be invoked afresh every time. This could be expensive. Can this be fixed? Yes, you can instead build a model that's input-invariant.

I'll send that along shortly.

[okhat]

class PrompterNLI(dspy.Signature):
    """
    Prepare a prompt with {premise} and {hypothesis} for an NLI BERT model.
    Use one [CLS] and two [SEP] tokens. Introduce creative prefixes or suffixes for the variables.
    """

    variables: list[str] = dspy.InputField()
    prompt: str = dspy.OutputField(desc="a pattern on which we may call .format()")


class PromptedBERT(dspy.Module):
    def __init__(self, bert_nli_classifier: Callable[[str], NLI]):
        self.bert_nli_classifier = bert_nli_classifier
        self.prompter = dspy.ChainOfThought(PrompterNLI, temperature=1.0)
    
    def forward(self, premise: str, hypothesis: str) -> NLI:
        prompt = self.prompter(variables=['premise', 'hypothesis']).prompt
        prompt = prompt.strip('"').format(premise=premise, hypothesis=hypothesis)
        return self.bert_nli_classifier(prompt)

Here, self.prompter always receives the same inputs, which means you can cache it after optimization.

Here's one sample output before optimization:

[CLS] Given the statement: {premise} [SEP] Assess the claim: {hypothesis} [SEP] Determine if the claim logically follows from the statement.

We've never done this before, though, so who knows how BERT will respond to these prompts!

[umarbutler]

Shukran Omar, this was a great starting point for me.

It might be helpful for me to give you a bit more detail about what I'm trying to do.

In particular, I have a zero-shot NLI classifier (eg, MoritzLaurer/deberta-v3-large-zeroshot-v2.0) that takes a piece of text (aka a premise) and a statement about that text (aka a hypothesis) and outputs a binary classification of whether or not the statement is supported by the text (aka entailment).

Sometimes, I might already have some labelled data which I can then use to optimise my statement.

For example, suppose I want to classify all of my blog posts by whether or not they relate to law. I could use a statement like 'This relates to law' but perhaps a better statement would be 'This blog post relates to law' or 'This blog post relates to the topic of law'. If I already have labelled data, I can simply try a bunch of different permutations of the same statement until I arrive at a statement that achieves the best Matthews' correlation coefficent or accuracy.

As you might imagine, this process can become quite involved if you need to repeat it for many different classification problems, models or hyperparameters.

I had a crack at adapting your code for this problem, however, it doesn't seem to be working because my training examples are missing a statement input. Is it possible to have DSPy come up with statements on its own without any examples of what an optimal statement looks like, solely by optimising for accuracy on the training set?

import dspy
import torch
import dspy.evaluate.metrics

from transformers import pipeline


class ZeroShotClassification(dspy.Signature):
    statement: str = dspy.InputField()
    text: str = dspy.InputField()
    supported: str = dspy.OutputField()

class ZeroShotClassifier(dspy.Module):
    def __init__(self, model_name: str, device: str = 'cuda' if torch.cuda.is_available() else 'cpu') -> None:
        self.pipe = pipeline('zero-shot-classification', model = model_name, device = device)
        self.prompter = dspy.ChainOfThought(ZeroShotClassification, temperature = 1.)
    
    def forward(self, text: str) -> str:
        statement = self.prompter(text = text).statement
        statement = statement.replace("'", '').replace('"', '').strip()
        
        return ['not supported', 'supported'][self.pipe(text, statement, hypothesis_template = '{}')['score'] > 0.5]

# BEGIN CONFIG #
ZERO_SHOT_CLASSIFIER_NAME = 'MoritzLaurer/deberta-v3-large-zeroshot-v2.0'
LM_NAME = 'openai/gpt-4o-mini'
THREADS = 1
AUTO_MODE = 'medium'
EXPERIMENTAL = True
NEEDS_PERMISSION = False
DEVICE = 'cpu'

SENTIMENT_EXAMPLES = [ # The goal would be to have DSPy come up with a prompt like 'This text expresses positive sentiment towards DSPy.'
    dspy.Example(
        text = 'I love DSPy!',
        supported = 'supported',
    ),
    dspy.Example(
        text = 'I hate DSPy!',
        supported = 'not supported',
    ),
    dspy.Example(
        text = "I'm not sure how I feel about DSPy.",
        supported = 'not supported',
    ),
]
# END CONFIG #

SENTIMENT_EXAMPLES = [example.with_inputs('text') for example in SENTIMENT_EXAMPLES]

lm = dspy.LM(LM_NAME)
dspy.configure(lm = lm, experimental = EXPERIMENTAL)
zsc = ZeroShotClassifier(ZERO_SHOT_CLASSIFIER_NAME, DEVICE)
optimizer = dspy.MIPROv2(metric = dspy.evaluate.metrics.answer_exact_match, auto = AUTO_MODE, num_threads = THREADS)
optimized_zsc = optimizer.compile(zsc, trainset = SENTIMENT_EXAMPLES, requires_permission_to_run = NEEDS_PERMISSION)

[chrico-bu-uab]

@umarbutler Maybe not perfect, but I think something like this is what you want?

import dspy
import dspy.evaluate.metrics
from dspy.evaluate.confusion import Confusion

from transformers import pipeline
from datasets import load_dataset

# BEGIN CONFIG #
ZERO_SHOT_CLASSIFIER_NAME = 'MoritzLaurer/deberta-v3-large-zeroshot-v2.0'
LM_NAME = 'openai/gpt-4.1-mini'
THREADS = 4
EXPERIMENTAL = True
NEEDS_PERMISSION = False
DEVICE = 'mps'

STATEMENT = 'The sentiment is neutral.'

EXAMPLES = load_dataset("mteb/tweet_sentiment_extraction")
FUNC = lambda example: "supported" if example["label"] == 1 else "not supported"
# END CONFIG #


class GenerateHypothesis(dspy.Signature):
    premise: str = dspy.InputField()
    hypothesis: str = dspy.OutputField()


class CompareHypotheses(dspy.Signature):
    hypothesis_a: str = dspy.InputField()
    hypothesis_b: str = dspy.InputField()
    are_aligned: bool = dspy.OutputField()


class ZeroShotClassifier(dspy.Module):
    def __init__(self, model_name: str, device: str, categories: list[str]) -> None:
        self.pipe = pipeline('zero-shot-classification', model=model_name, device=device)
        self.generator = dspy.ChainOfThought(GenerateHypothesis, temperature=1.)
        self.comparer = dspy.ChainOfThought(CompareHypotheses)
        self.threshold = 0.5
        self.categories = categories

    def forward(self, actual: str, text: str):
        statement = self.generator(premise=text).hypothesis.replace("'", '').replace('"', '').strip()
        compared = self.comparer(hypothesis_a=statement, hypothesis_b=actual).are_aligned
        if self.pipe(text, statement, hypothesis_template='{}')['scores'][0] < self.threshold:
            compared = not compared
        return dspy.Prediction(answer=self.categories[compared])

EXAMPLES = [[dspy.Example(actual=STATEMENT, text=example["text"], answer=FUNC(example),
                          ) for example in EXAMPLES[x]] for x in ["train", "test"]]
train, test = [[example.with_inputs('actual', 'text') for example in x] for x in EXAMPLES]
train = train[:100]
test = test[:100]

lm = dspy.LM(LM_NAME)
dspy.configure(lm=lm, experimental=EXPERIMENTAL)
zsc = ZeroShotClassifier(ZERO_SHOT_CLASSIFIER_NAME, DEVICE, ["not supported", "supported"])
evaluate = Confusion(devset=test,
                      num_threads=THREADS,
                     display_progress=True,
                     return_matrix=True,
                     return_outputs=True,
                     output_field="answer",
                     provide_traceback=True,
                     max_errors=0)
print(evaluate(zsc))
optimizer = dspy.MIPROv2(metric=dspy.evaluate.answer_exact_match, num_threads=THREADS, auto="heavy")
optimized_zsc = optimizer.compile(zsc, trainset=train, requires_permission_to_run=NEEDS_PERMISSION)
optimized_zsc.save('zsc.json')
print(evaluate(optimized_zsc))

Optimizing brought accuracy from 51% to 68% and MCC from 0 to 0.33!

{
  "generator.predict": {
    "traces": [],
    "train": [],
    "demos": [],
    "signature": {
      "instructions": "Given a short informal text labeled as neutral sentiment (the `premise`), carefully analyze the text's emotional cues, contextual subtleties, slang, emoticons, and any ambiguous expressions typical of conversational or social media language. Then, reason step-by-step to generate a clear, logical hypothesis that accurately reflects the underlying sentiment or implication of the text, verifying whether the neutral label is supported or contradicted by the emotional and pragmatic context. Your output should include a detailed reasoning process that explains how the premise leads to the hypothesis, followed by the hypothesis itself, which should be a concise statement summarizing the inferred sentiment or stance implied by the premise.",
      "fields": [
        {
          "prefix": "Premise:",
          "description": "${premise}"
        },
        {
          "prefix": "Reasoning: Let's think step by step in order to",
          "description": "${reasoning}"
        },
        {
          "prefix": "Hypothesis:",
          "description": "${hypothesis}"
        }
      ]
    },
    "lm": null
  },
  "comparer.predict": {
    "traces": [],
    "train": [],
    "demos": [
      {
        "augmented": true,
        "hypothesis_a": "I am sad because I will no longer see you in San Diego.",
        "hypothesis_b": "The sentiment is neutral.",
        "reasoning": "Hypothesis A expresses sadness about not seeing someone in San Diego, which clearly conveys a negative emotional sentiment. Hypothesis B states that the sentiment is neutral, which contradicts the expressed sadness in Hypothesis A. Therefore, the two hypotheses are not aligned in terms of sentiment.",
        "are_aligned": false
      },
      {
        "augmented": true,
        "hypothesis_a": "I feel stressed at work because of how my boss treats me.",
        "hypothesis_b": "The sentiment is neutral.",
        "reasoning": "Hypothesis A expresses a negative sentiment, specifically stress caused by the boss's treatment. Hypothesis B states that the sentiment is neutral, which contradicts the negative sentiment expressed in Hypothesis A. Therefore, the two hypotheses are not aligned.",
        "are_aligned": false
      },
      {
        "augmented": true,
        "hypothesis_a": "The person will try to attend the music event tonight despite having lost their voice.",
        "hypothesis_b": "The sentiment is neutral.",
        "reasoning": "Hypothesis A describes a specific scenario about a person attempting to attend a music event despite losing their voice, which implies a certain determination or effort. Hypothesis B states that the sentiment is neutral, which is a general emotional tone assessment rather than a specific event or action. Since Hypothesis A is a detailed statement about behavior and Hypothesis B is a broad sentiment classification, they do not align in content or intent.",
        "are_aligned": false
      },
      {
        "augmented": true,
        "hypothesis_a": "The baby brings joy during late-night feedings.",
        "hypothesis_b": "The sentiment is neutral.",
        "reasoning": "Hypothesis A expresses a positive sentiment by stating that the baby brings joy during late-night feedings, which implies happiness or pleasure. Hypothesis B states that the sentiment is neutral, which contradicts the positive sentiment expressed in Hypothesis A. Therefore, the two hypotheses are not aligned in terms of sentiment.",
        "are_aligned": false
      }
    ],
    "signature": {
      "instructions": "Given two hypotheses, `hypothesis_a` and `hypothesis_b`, carefully analyze and compare their semantic content, emotional cues, and contextual implications. Produce a detailed, step-by-step chain-of-thought reasoning that explains whether the two hypotheses convey consistent or aligned meanings, especially focusing on the presence or absence of emotional sentiment and neutrality. Finally, determine and output a boolean field `are_aligned` indicating if the hypotheses are aligned (True) or not (False) based on this thorough reasoning. Your explanation should clarify how the emotional tone and contextual subtleties in `hypothesis_a` either support or contradict the sentiment expressed in `hypothesis_b`.",
      "fields": [
        {
          "prefix": "Hypothesis A:",
          "description": "${hypothesis_a}"
        },
        {
          "prefix": "Hypothesis B:",
          "description": "${hypothesis_b}"
        },
        {
          "prefix": "Reasoning: Let's think step by step in order to",
          "description": "${reasoning}"
        },
        {
          "prefix": "Are Aligned:",
          "description": "${are_aligned}"
        }
      ]
    },
    "lm": null
  },
  "metadata": {
    "dependency_versions": {
      "python": "3.9",
      "dspy": "2.6.23",
      "cloudpickle": "3.1"
    }
  }
}

(Confusion is a custom class I made for calculating MCC, you can just use dspy.Evaluate and accuracy.)