Get your LLM to pick the right category.
Project description
CAPPr: Completion After Prompt Probability
CAPPr performs text classification. No training. No post-processing. Just have your LLM pick from a list of choices. Or compute the probability of a completion given a prompt. Squeeze more out of open source LLMs.
Usage
Use a model from the OpenAI API
Specifically, this model must be compatible with the /v1/completions endpoint.
from cappr.openai.classify import predict
prompt = """
Tweet about a movie: "Oppenheimer was pretty good. But 3 hrs...cmon Nolan."
This tweet contains the following criticism:
""".strip("\n")
completions = ("bad message", "too long", "unfunny")
pred = predict(prompt, completions, model="text-ada-001")
print(pred)
# 'too long'
Notice that the completions can contain many tokens.
Extract the final answer from a step-by-step completion
Step-by-step and chain-of-thought prompts are highly effective ways to get an LLM to "reason" about more complex tasks. But if you need a structured output, a step-by-step completion is unwieldy. Use CAPPr to extract the final answer from these types of completions, given a list of possible answers.
See this idea in action here in the docs. CAPPr is 100% guaranteed to return an output from the list of answers.
Use a model from the HuggingFace model hub
Specifically, this model must be able to be loaded using
transformers.AutoModelForCausalLM.from_pretrained
.
from transformers import AutoModelForCausalLM, AutoTokenizer
from cappr.huggingface.classify import predict
# Load a model and its corresponding tokenizer
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
prompt = "Which planet is closer to the Sun: Mercury or Earth?"
completions = ("Mercury", "Earth")
pred = predict(prompt, completions, model_and_tokenizer=(model, tokenizer))
print(pred)
# 'Mercury'
For an example with Llama 2, see the notebook
demos/llama2/copa.ipynb
.
So far, CAPPr has been tested for correctness on the following architectures:
- GPT-2
- GPT-J
- GPT-NeoX (including StableLM, and its tuned/instruct and GPTQd versions)
- Llama
- Llama 2 (chat, raw, and its GPTQd versions)
- Mistral.
Raise an issue to lmk that you don't see your architecture on this list.
Use an AutoGPTQ model
cappr.huggingface
seems to work fine with models loaded via
auto_gptq.AutoGPTQForCausalLM.from_quantized
.
But I haven't thoroughly tested that. See this
notebook for a demo.
Run in batches
Also, let's predict probabilities instead of the class.
from transformers import AutoModelForCausalLM, AutoTokenizer
from cappr.huggingface.classify import predict_proba
# Load a model and its corresponding tokenizer
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
prompts = [
"Stephen Curry is a",
"Martina Navratilova was a",
"Dexter, from the TV Series Dexter's Laboratory, is a",
"LeBron James is a",
]
# Each of the prompts could be completed with one of these:
class_names = ("basketball player", "tennis player", "scientist")
prior = ( 1/6, 1/6, 2/3 )
# Say I expect most of my data to have scientists
# Run CAPPr
pred_probs = predict_proba(
prompts=prompts,
completions=class_names,
model_and_tokenizer=(model, tokenizer),
batch_size=32, # whatever fits on your CPU/GPU
prior=prior,
)
# pred_probs[i,j] = probability that prompts[i] is classified as class_names[j]
print(pred_probs.round(1))
# [[0.5 0.3 0.2]
# [0.3 0.6 0.2]
# [0.1 0.1 0.8]
# [0.8 0.2 0. ]]
# For each prompt, which completion is most likely?
pred_class_idxs = pred_probs.argmax(axis=-1)
preds = [class_names[pred_class_idx] for pred_class_idx in pred_class_idxs]
print(preds)
# ['basketball player',
# 'tennis player',
# 'scientist',
# 'basketball player']
Run in batches, where each prompt has a different set of possible completions
Again, let's use cappr.huggingface
to predict probabilities.
from transformers import AutoModelForCausalLM, AutoTokenizer
from cappr.huggingface.classify import predict_proba_examples
from cappr import Example
# Load a model and its corresponding tokenizer
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Create a sequence of Example objects representing your classification tasks
examples = [
Example(
prompt="Jodie Foster played",
completions=("Clarice Starling", "Trinity in The Matrix"),
),
Example(
prompt="Batman, from Batman: The Animated Series, was played by",
completions=("Pete Holmes", "Kevin Conroy", "Spongebob!"),
prior= ( 1/3 , 2/3 , 0 ),
),
]
# Run CAPPr
pred_probs = predict_proba_examples(examples, model_and_tokenizer=(model, tokenizer))
# pred_probs[i][j] = probability that examples[i].prompt is classified as
# examples[i].completions[j]
print([example_pred_probs.round(2) for example_pred_probs in pred_probs])
# [array([0.7, 0.3]),
# array([0.03, 0.97, 0. ])]
# For each example, which completion is most likely?
pred_class_idxs = [example_pred_probs.argmax() for example_pred_probs in pred_probs]
preds = [
example.completions[pred_class_idx]
for example, pred_class_idx in zip(examples, pred_class_idxs)
]
print(preds)
# ['Clarice Starling',
# 'Kevin Conroy']
See
demos/llama2/copa.ipynb
for a demonstration of a slightly harder classification task.
Documentation
Installation
To use HuggingFace models:
pip install "cappr[hf]"
Or, to use OpenAI models, sign up for the OpenAI
API, set the environment variable OPENAI_API_KEY
, and then
pip install cappr
(Optional) Install requirements for running the repo's
demos
pip install "cappr[demos]"
Motivation
Minimize engineering complexity.
Current approaches, based on generating text, require the following workflow:
- Design a prompt which asks the model to output exactly one choice from a given set of choices
- Given this prompt, figure out how often the model doesn't output one of the given choices—call this an "invalid" output
- Figure out how to post-processes invalid outputs, depending on the way they look. Or give up, and figure out how to get your application to gracefully fail on an invalid output
- Figure out if you need to tweak the text generation strategy, loop back to step (2).
This workflow is not scalable, and not necessary. CAPPr is guaranteed to output exactly one choice from a given set of choices. As a result, your work is reduced to designing a prompt-completion string format.
Long
Cool
A handful of experiments suggest that CAPPr squeezes more out of smaller LLMs. See this page of the documentation.
Honest
am bored. am unemployed.
Performance
Statistical performance
Not too shabby. TODO: summary table comparing competing methods.
For HuggingFace models, see
- the 4 GB Llama 2 COPA demo
- and this (minimal but surprising) 3 GB StableLM demo.
I'll evaluate Llama 2 or Mistral 7B on a few more datasets.
Computational performance
How it works
You input a prompt
string, a end_of_prompt
string (a whitespace or empty) and a set
of candidate completion
strings such that the string—
{prompt}{end_of_prompt}{completion}
—is a naturally flowing thought. CAPPr picks the completion
which is mostly likely to
follow prompt
by computing the:
Completion
After
Prompt
Probability
The method is fleshed out in my question on Cross Validated.
Related work
See this page of the documentation.
Local development
Setup
-
Create a new Python 3.8+ virtual environment. Activate the venv. I use
virtualenvwrapper
. For this example, let's create a virtual environment calledcappr
using Python's nativevenv
:cd your/venvs python3 -m venv cappr source cappr/bin/activate python -m pip install wheel --upgrade pip
-
cd
to wherever you store projects, and clone the repo (or fork it and clone that) therecd your/projects git clone https://github.com/kddubey/cappr.git
-
cd
to the repo and install this package in editable mode, along with development requirements (after ensuring that your venv is activated!)cd cappr python -m pip install -e ".[dev]"
VS code extensions for development
- autoDocstring. Use the numpy format, and check "Start On New Line".
- Set Python formatting to
black
. - Rewrap. Enable Auto Wrap.
And set the vertical line ruler to 88.
Testing
From the repo home directory cappr
:
pytest
Note that a few small transformers and tokenizers will be downloaded to your computer.
Sometimes I get paranoid about bigger code changes. So consider additionally testing
statistical performance by running an appropriate demo in
demos
.
Docs
To test changes to documentation, first locally build them from the repo home directory
cappr
via
cd docs
make html
and then preview them by opening docs/build/html/index.html
in your browser.
After merging code to main, the official docs will be automatically built and published.
Release
Bump the version, and then create a new release on GitHub. A new version of the package will then be automatically published on PyPI.
Todo
I'm dumping TODOs here:
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.