Few-shot¶

Few-shot learning is a simple, yet surprisingly effective, input control method for steering a language model's behavior by including examples of desirable/undesirable behavior in the prompt (Brown et al., 2020; Zhao et al., 2021). This notebook illustrates how few-shot learning is implemented in the toolkit (via the FewShot class). The toolkit contains few-shot steering under the following two modes:

1. Runtime Examples Mode: Passes specific examples directly at generation time via runtime_kwargs.

2. Pool Sampling Mode: Defines (positive and negative) example pools during initialization and a strategy for choosing a specified number of examples from the pools at runtime.

In the pool sampling mode, the specific examples drawn from the pool at each call are determined by a selector. The default selector (RandomSelector) samples uniformly from each polarity pool, which is appropriate when the pool is small and homogeneous. To enable more control for larger or more heterogeneous pools, the few-shot implementation also allows for the definition of more complex selectors. One such selector is a learned dense retriever via EPRSelector (Rubin, Herzig, Berant 2021) that ranks pool items by their similarity to the prompt. The EPR selector is illustrated at the end of this notebook.

Selectors are configured via the selector argument of FewShot. Pass a string name (e.g., "random") for one of the built-in selectors, or an instance of a BaseSelector subclass for a custom or pre-configured selector. The selector argument only applies to Pool Sampling Mode; in Runtime Examples Mode the examples are used exactly as passed in.

In this demo, we'll show how FewShot can be used to steer a model to respond more concisely.

Method parameters¶

parameter	type	description
selector	BaseSelector \| str \| None	Selector for picking examples from the pool. Accepts a string name ("random") or a BaseSelector instance (e.g. EPRSelector for learned retrieval). Defaults to random selection.
formatter	BaseFormatter \| None	Formatter that renders the example block into the chat / token stream. Defaults to FewShotBlockFormatter().
directive	str \| None	Directive statement at the beginning of the system prompt.
positive_example_pool	list[dict] \| None	Pool of positive examples to sample from at runtime.
negative_example_pool	list[dict] \| None	Pool of negative examples to sample from at runtime.
k_positive	int \| None	Number of positive examples to sample from the pool. Required if positive_example_pool is provided.
k_negative	int \| None	Number of negative examples to sample from the pool. Required if negative_example_pool is provided.

Setup¶

If running this from a Google Colab notebook, please uncomment the following cell to install the toolkit. The following block is not necessary if running this notebook from a virtual environment where the package has already been installed.

In [1]:

# !git clone https://github.com/generative-computing/steerability.git
# %cd steerability

# !git clone https://github.com/generative-computing/steerability.git # %cd steerability

The following authentication steps may be necessary to access any gated models (after being granted access by Hugging Face). Uncomment the following if you need to log in to the Hugging Face Hub using your token stored in the .env file:

In [2]:

# !pip install -q python-dotenv
# from dotenv import load_dotenv
# import os

# load_dotenv()
# token = os.getenv("HUGGINGFACE_TOKEN")
# from huggingface_hub import login
# login(token=token)

# !pip install -q python-dotenv # from dotenv import load_dotenv # import os # load_dotenv() # token = os.getenv("HUGGINGFACE_TOKEN") # from huggingface_hub import login # login(token=token)

Example: Steering for conciseness¶

In [3]:

from transformers import AutoModelForCausalLM, AutoTokenizer

from steerability.algorithms.input_control.few_shot.control import FewShot
from steerability.algorithms.core.steering_pipeline import SteeringPipeline

MODEL_NAME = "google/gemma-3-4b-it"

from transformers import AutoModelForCausalLM, AutoTokenizer from steerability.algorithms.input_control.few_shot.control import FewShot from steerability.algorithms.core.steering_pipeline import SteeringPipeline MODEL_NAME = "google/gemma-3-4b-it"

The following example illustrates how to steer a model's behavior to respond more concisely. We've defined some positive examples (those that represent the desired behavior) and negative examples (those that represent the undesired behavior) below.

In [4]:

positive_examples = [
    {"question": "What's the capital of France?", "answer": "Paris"},
    {"question": "How many miles is it to the moon?", "answer": "238,855"},
    {"question": "What's the boiling point of water?", "answer": "100°C"},
    {"question": "How many days in a leap year?", "answer": "366"},
    {"question": "What's the speed of light?", "answer": "299,792,458 m/s"},
    {"question": "What's 15% of 200?", "answer": "30"},
    {"question": "How many continents are there?", "answer": "7"},
    {"question": "What's the atomic number of gold?", "answer": "79"}
]

negative_examples = [
    {"question": "What's the capital of France?", "answer": "The capital of France is Paris. Located along the Seine River in the country's north-central region, it is the political, cultural, and economic heart of the nation."},
    {"question": "How many miles is it to the moon?", "answer": "The Moon sits about 238,855 miles from Earth on average. Of course, because its orbit is elliptical, that distance ranges from roughly 226,000 miles at perigee to 252,000 miles at apogee."},
    {"question": "What's the boiling point of water?", "answer": "You're asking about the boiling point of water. The boiling point of water is 100 degrees Celsius at sea level."},
    {"question": "How many days in a leap year?", "answer": "A leap year contains 366 days. This is one more than a standard 365-day year, with the extra day added to February (the 29th) to keep the calendar aligned with Earth's orbit around the Sun."},
    {"question": "What's the speed of light?", "answer": "The speed of light in a vacuum is approximately 299,792,458 meters per second, though in everyday contexts it's often rounded to about 300,000 kilometers per second for convenience."},
    {"question": "What's 15% of 200?", "answer": "To find 15% of 200, multiply 200 by 0.15. That gives 200 × 0.15 = 30, so the answer is 30."},
    {"question": "How many continents are there?", "answer": "There are 7 continents on Earth:\n\n1. Africa\n2. Antarctica\n3. Asia\n4. Europe\n5. North America\n6. Oceania\n7. South America"},
    {"question": "What's the atomic number of gold?", "answer": "Gold's atomic number is 79. Its chemical symbol, Au, derives from the Latin word 'aurum,' meaning 'shining dawn.'"}
]

positive_examples = [ {"question": "What's the capital of France?", "answer": "Paris"}, {"question": "How many miles is it to the moon?", "answer": "238,855"}, {"question": "What's the boiling point of water?", "answer": "100°C"}, {"question": "How many days in a leap year?", "answer": "366"}, {"question": "What's the speed of light?", "answer": "299,792,458 m/s"}, {"question": "What's 15% of 200?", "answer": "30"}, {"question": "How many continents are there?", "answer": "7"}, {"question": "What's the atomic number of gold?", "answer": "79"} ] negative_examples = [ {"question": "What's the capital of France?", "answer": "The capital of France is Paris. Located along the Seine River in the country's north-central region, it is the political, cultural, and economic heart of the nation."}, {"question": "How many miles is it to the moon?", "answer": "The Moon sits about 238,855 miles from Earth on average. Of course, because its orbit is elliptical, that distance ranges from roughly 226,000 miles at perigee to 252,000 miles at apogee."}, {"question": "What's the boiling point of water?", "answer": "You're asking about the boiling point of water. The boiling point of water is 100 degrees Celsius at sea level."}, {"question": "How many days in a leap year?", "answer": "A leap year contains 366 days. This is one more than a standard 365-day year, with the extra day added to February (the 29th) to keep the calendar aligned with Earth's orbit around the Sun."}, {"question": "What's the speed of light?", "answer": "The speed of light in a vacuum is approximately 299,792,458 meters per second, though in everyday contexts it's often rounded to about 300,000 kilometers per second for convenience."}, {"question": "What's 15% of 200?", "answer": "To find 15% of 200, multiply 200 by 0.15. That gives 200 × 0.15 = 30, so the answer is 30."}, {"question": "How many continents are there?", "answer": "There are 7 continents on Earth:\n\n1. Africa\n2. Antarctica\n3. Asia\n4. Europe\n5. North America\n6. Oceania\n7. South America"}, {"question": "What's the atomic number of gold?", "answer": "Gold's atomic number is 79. Its chemical symbol, Au, derives from the Latin word 'aurum,' meaning 'shining dawn.'"} ]

To analyze the model's conciseness, we define a prompt (question) that asks a question that can admit a concise answer:

In [5]:

PROMPT = "How many ounces are in a pint?"

PROMPT = "How many ounces are in a pint?"

Baseline model behavior¶

The baseline model behavior is generated by simply passing the templated and tokenized prompt into the base model's generate.

In [6]:

model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

chat = tokenizer.apply_chat_template(
    [{"role": "user", "content": PROMPT}], 
    tokenize=False, 
    add_generation_prompt=True
)
inputs = tokenizer(chat, return_tensors="pt")

baseline_outputs = model.generate(
    **inputs.to(model.device), 
    max_new_tokens=150
)

print("\nResponse (baseline):\n")
print(tokenizer.decode(baseline_outputs[0][len(inputs['input_ids'][0]):], skip_special_tokens=True))

model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto") tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) chat = tokenizer.apply_chat_template( [{"role": "user", "content": PROMPT}], tokenize=False, add_generation_prompt=True ) inputs = tokenizer(chat, return_tensors="pt") baseline_outputs = model.generate( **inputs.to(model.device), max_new_tokens=150 ) print("\nResponse (baseline):\n") print(tokenizer.decode(baseline_outputs[0][len(inputs['input_ids'][0]):], skip_special_tokens=True))

Response (baseline):

There are **16 ounces** in a pint. 

It's important to note that there are two different sizes of a "pint" depending on the context:

*   **U.S. Pint:** 16 fluid ounces (commonly used for beverages like beer or milk)
*   **Imperial Pint:** 20 fluid ounces (used in the UK and some other countries)

Since you asked for the number of ounces, I'm assuming you mean the U.S. pint (16 ounces).

Steering via runtime_kwargs¶

In this mode, examples are passed in at generate() time. This allows for control over which examples are used for each generation. Note that the instantiation of FewShot in this mode does not require any arguments.

In [7]:

few_shot_runtime = FewShot()

few_shot_runtime = FewShot()

Given the control, we define the steering pipeline (via SteeringPipeline) and steer it (performs some lightweight initialization of FewShot):

In [8]:

few_shot_runtime_pipeline = SteeringPipeline(
    model_name_or_path=MODEL_NAME,
    controls=[few_shot_runtime],
    device_map="auto"
)
few_shot_runtime_pipeline.steer()

few_shot_runtime_pipeline = SteeringPipeline( model_name_or_path=MODEL_NAME, controls=[few_shot_runtime], device_map="auto" ) few_shot_runtime_pipeline.steer()

Inference on the steered model can then be run as usual to generate the steered output. Note the specific examples listed above are passed directly into generate via the runtime_kwargs argument.

In [9]:

inputs = tokenizer(PROMPT, return_tensors="pt")

output = few_shot_runtime_pipeline.generate(
    input_ids=inputs.input_ids,
    runtime_kwargs={
        "positive_examples": positive_examples,
        "negative_examples": negative_examples
    },
    max_new_tokens=50,
    temperature=0.7,
    return_full_sequence=False
)

print("\nResponse (FewShot w/ fixed examples):\n")
print(few_shot_runtime_pipeline.tokenizer.decode(output[0], skip_special_tokens=True))

inputs = tokenizer(PROMPT, return_tensors="pt") output = few_shot_runtime_pipeline.generate( input_ids=inputs.input_ids, runtime_kwargs={ "positive_examples": positive_examples, "negative_examples": negative_examples }, max_new_tokens=50, temperature=0.7, return_full_sequence=False ) print("\nResponse (FewShot w/ fixed examples):\n") print(few_shot_runtime_pipeline.tokenizer.decode(output[0], skip_special_tokens=True))

/dccstor/principled_ai/users/erikmiehling/AISteer360/steerability/algorithms/core/utils/controls.py:163: UserWarning: FewShot override(s) `adapt_messages` but received tensor/text input; the message-level adaptation will not run. Pass `list[dict]` or `list[list[dict]]` to engage `adapt_messages`.
  warnings.warn(

Response (FewShot w/ fixed examples):

16

Steering via example pools and a selector¶

In some cases, the requirement to pass in specific examples may not be necessary or even desirable, e.g., if you have a large pool of examples and are not sure which yield the desired behavior. To accommodate this, we allow for the user to specify example pools and a selector for how to sample from the pool.

First, clear the memory from the previous mode:

In [10]:

import gc, torch
gc.collect()
torch.cuda.empty_cache()

import gc, torch gc.collect() torch.cuda.empty_cache()

In [11]:

positive_example_pool = [
    {"question": "What's the capital of France?", "answer": "Paris"},
    {"question": "How many miles is it to the moon?", "answer": "238,855"},
    {"question": "What's the boiling point of water?", "answer": "100°C"},
    {"question": "How many days in a leap year?", "answer": "366"},
    {"question": "What's the speed of light?", "answer": "299,792,458 m/s"},
    {"question": "What's 15% of 200?", "answer": "30"},
    {"question": "How many continents are there?", "answer": "7"},
    {"question": "What's the atomic number of gold?", "answer": "79"},
    {"question": "What's the capital of Japan?", "answer": "Tokyo"},
    {"question": "How many sides does a hexagon have?", "answer": "6"},
    {"question": "What's 9 * 7?", "answer": "63"},
    {"question": "What's the freezing point of water?", "answer": "0°C"},
    {"question": "How many planets are in the Solar System?", "answer": "8"},
    {"question": "What's the chemical symbol for sodium?", "answer": "Na"},
    {"question": "What's the largest ocean on Earth?", "answer": "Pacific Ocean"},
    {"question": "How many degrees are in a right angle?", "answer": "90"},
    {"question": "What's the square root of 144?", "answer": "12"},
    {"question": "Who's the author of '1984'?", "answer": "George Orwell"},
    {"question": "What's the currency of the United Kingdom?", "answer": "Pound sterling"},
    {"question": "What gas do plants primarily absorb during photosynthesis?", "answer": "Carbon dioxide"},
    {"question": "How many letters are in the English alphabet?", "answer": "26"},
    {"question": "What's the largest planet in our solar system?", "answer": "Jupiter"},
    {"question": "What's the tallest mountain in the world?", "answer": "Mount Everest"},
    {"question": "What's the primary language spoken in Brazil?", "answer": "Portuguese"},
    {"question": "What is the Roman numeral for 50?", "answer": "L"},
    {"question": "How many hours are in two days?", "answer": "48"},
    {"question": "What's 3/4 as a percentage?", "answer": "75%"},
    {"question": "What's the chemical formula for table salt?", "answer": "NaCl"},
    {"question": "How many bits are in a byte?", "answer": "8"},
    {"question": "What's the smallest prime number?", "answer": "2"},
    {"question": "What's Pi rounded to two decimal places?", "answer": "3.14"},
    {"question": "How many bones are in the adult human body?", "answer": "206"}
]

negative_example_pool = [
    {"question": "What's the capital of France?", "answer": "The capital of France is Paris. Located along the Seine River in the country's north-central region, it is the political, cultural, and economic heart of the nation."},
    {"question": "How many miles is it to the moon?", "answer": "The Moon sits about 238,855 miles from Earth on average. Of course, because its orbit is elliptical, that distance ranges from roughly 226,000 miles at perigee to 252,000 miles at apogee."},
    {"question": "What's the boiling point of water?", "answer": "You're asking about the boiling point of water. The boiling point of water is 100 degrees Celsius at sea level."},
    {"question": "How many days in a leap year?", "answer": "A leap year contains 366 days. This is one more than a standard 365-day year, with the extra day added to February (the 29th) to keep the calendar aligned with Earth's orbit around the Sun."},
    {"question": "What's the speed of light?", "answer": "The speed of light in a vacuum is approximately 299,792,458 meters per second, though in everyday contexts it's often rounded to about 300,000 kilometers per second for convenience."},
    {"question": "What's 15% of 200?", "answer": "To find 15% of 200, multiply 200 by 0.15. That gives 200 × 0.15 = 30, so the answer is 30."},
    {"question": "How many continents are there?", "answer": "There are 7 continents on Earth:\n\n1. Africa\n2. Antarctica\n3. Asia\n4. Europe\n5. North America\n6. Oceania\n7. South America"},
    {"question": "What's the atomic number of gold?", "answer": "Gold's atomic number is 79. Its chemical symbol, Au, derives from the Latin word 'aurum,' meaning 'shining dawn.'"},
    {"question": "What's the capital of Japan?", "answer": "Sure! The capital of Japan is Tokyo, which is also one of the largest metropolitan areas in the world."},
    {"question": "How many sides does a hexagon have?", "answer": "A hexagon has 6 sides. The prefix 'hex-' comes from the Greek word for six, which is where the name originates."},
    {"question": "What's 9 * 7?", "answer": "You'd like to know what 9 times 7 is. Nine multiplied by seven equals 63."},
    {"question": "What's the freezing point of water?", "answer": "Water freezes at 0 degrees Celsius (32 degrees Fahrenheit) under normal conditions, although factors like pressure and dissolved impurities can lower this point somewhat."},
    {"question": "How many planets are in the Solar System?", "answer": "There are 8 planets in our Solar System. Feel free to ask if you'd like to know more about any of them!"},
    {"question": "What's the chemical symbol for sodium?", "answer": "**Sodium**\n\nThe chemical symbol for sodium is **Na**, which comes from its Latin name, *natrium*."},
    {"question": "What's the largest ocean on Earth?", "answer": "The largest ocean on Earth is the Pacific Ocean. It covers more than 60 million square miles — larger than all of Earth's landmasses combined."},
    {"question": "How many degrees are in a right angle?", "answer": "A right angle measures 90 degrees. It's the angle you'd see in the corner of a square, formed when two lines meet perpendicular to one another."},
    {"question": "What's the square root of 144?", "answer": "To find the square root of 144, we look for the number that, multiplied by itself, equals 144. Since 12 × 12 = 144, the square root of 144 is 12."},
    {"question": "Who's the author of '1984'?", "answer": "'1984' was written by George Orwell, the pen name of Eric Arthur Blair. It was published in 1949 and remains one of the most influential dystopian novels ever written."},
    {"question": "What's the currency of the United Kingdom?", "answer": "Great question! The currency of the United Kingdom is the pound sterling, often symbolized by £."},
    {"question": "What gas do plants primarily absorb during photosynthesis?", "answer": "During photosynthesis, plants primarily absorb carbon dioxide. They take it in through tiny pores called stomata and, using sunlight, convert it along with water into glucose and oxygen."},
    {"question": "How many letters are in the English alphabet?", "answer": "It's worth noting that the English alphabet contains 26 letters, ranging from A through Z."},
    {"question": "What's the largest planet in our solar system?", "answer": "The largest planet in our solar system is Jupiter, a gas giant so massive it could fit all the other planets inside it with room to spare."},
    {"question": "What's the tallest mountain in the world?", "answer": "The tallest mountain in the world is Mount Everest, at least when measured by height above sea level — by other measures, such as base-to-peak height, Mauna Kea would actually take the title."},
    {"question": "What's the primary language spoken in Brazil?", "answer": "The primary language spoken in Brazil is Portuguese. Let me know if you'd like to learn a few common phrases!"},
    {"question": "What is the Roman numeral for 50?", "answer": "The Roman numeral for 50 is L. The Romans used a system of letters to represent numbers, where L specifically stands for fifty."},
    {"question": "How many hours are in two days?", "answer": "Since one day has 24 hours, two days would be 24 × 2. That comes out to 48 hours."},
    {"question": "What's 3/4 as a percentage?", "answer": "To convert 3/4 to a percentage, divide 3 by 4 to get 0.75, then multiply by 100. That gives 75%."},
    {"question": "What's the chemical formula for table salt?", "answer": "The chemical formula for table salt is NaCl. It's an ionic compound made up of sodium (Na) and chlorine (Cl) atoms bonded together."},
    {"question": "How many bits are in a byte?", "answer": "Fun fact: there are 8 bits in a byte. This grouping is the standard unit used to represent a single character of text in most computing systems."},
    {"question": "What's the smallest prime number?", "answer": "The smallest prime number is 2. It's also the only even prime number, since every other even number is divisible by 2."},
    {"question": "What's Pi rounded to two decimal places?", "answer": "Sure thing! Pi rounded to two decimal places is 3.14."},
    {"question": "How many bones are in the adult human body?", "answer": "An adult human body has 206 bones, though the exact number can vary slightly from person to person depending on how certain small bones are counted."}
]

positive_example_pool = [ {"question": "What's the capital of France?", "answer": "Paris"}, {"question": "How many miles is it to the moon?", "answer": "238,855"}, {"question": "What's the boiling point of water?", "answer": "100°C"}, {"question": "How many days in a leap year?", "answer": "366"}, {"question": "What's the speed of light?", "answer": "299,792,458 m/s"}, {"question": "What's 15% of 200?", "answer": "30"}, {"question": "How many continents are there?", "answer": "7"}, {"question": "What's the atomic number of gold?", "answer": "79"}, {"question": "What's the capital of Japan?", "answer": "Tokyo"}, {"question": "How many sides does a hexagon have?", "answer": "6"}, {"question": "What's 9 * 7?", "answer": "63"}, {"question": "What's the freezing point of water?", "answer": "0°C"}, {"question": "How many planets are in the Solar System?", "answer": "8"}, {"question": "What's the chemical symbol for sodium?", "answer": "Na"}, {"question": "What's the largest ocean on Earth?", "answer": "Pacific Ocean"}, {"question": "How many degrees are in a right angle?", "answer": "90"}, {"question": "What's the square root of 144?", "answer": "12"}, {"question": "Who's the author of '1984'?", "answer": "George Orwell"}, {"question": "What's the currency of the United Kingdom?", "answer": "Pound sterling"}, {"question": "What gas do plants primarily absorb during photosynthesis?", "answer": "Carbon dioxide"}, {"question": "How many letters are in the English alphabet?", "answer": "26"}, {"question": "What's the largest planet in our solar system?", "answer": "Jupiter"}, {"question": "What's the tallest mountain in the world?", "answer": "Mount Everest"}, {"question": "What's the primary language spoken in Brazil?", "answer": "Portuguese"}, {"question": "What is the Roman numeral for 50?", "answer": "L"}, {"question": "How many hours are in two days?", "answer": "48"}, {"question": "What's 3/4 as a percentage?", "answer": "75%"}, {"question": "What's the chemical formula for table salt?", "answer": "NaCl"}, {"question": "How many bits are in a byte?", "answer": "8"}, {"question": "What's the smallest prime number?", "answer": "2"}, {"question": "What's Pi rounded to two decimal places?", "answer": "3.14"}, {"question": "How many bones are in the adult human body?", "answer": "206"} ] negative_example_pool = [ {"question": "What's the capital of France?", "answer": "The capital of France is Paris. Located along the Seine River in the country's north-central region, it is the political, cultural, and economic heart of the nation."}, {"question": "How many miles is it to the moon?", "answer": "The Moon sits about 238,855 miles from Earth on average. Of course, because its orbit is elliptical, that distance ranges from roughly 226,000 miles at perigee to 252,000 miles at apogee."}, {"question": "What's the boiling point of water?", "answer": "You're asking about the boiling point of water. The boiling point of water is 100 degrees Celsius at sea level."}, {"question": "How many days in a leap year?", "answer": "A leap year contains 366 days. This is one more than a standard 365-day year, with the extra day added to February (the 29th) to keep the calendar aligned with Earth's orbit around the Sun."}, {"question": "What's the speed of light?", "answer": "The speed of light in a vacuum is approximately 299,792,458 meters per second, though in everyday contexts it's often rounded to about 300,000 kilometers per second for convenience."}, {"question": "What's 15% of 200?", "answer": "To find 15% of 200, multiply 200 by 0.15. That gives 200 × 0.15 = 30, so the answer is 30."}, {"question": "How many continents are there?", "answer": "There are 7 continents on Earth:\n\n1. Africa\n2. Antarctica\n3. Asia\n4. Europe\n5. North America\n6. Oceania\n7. South America"}, {"question": "What's the atomic number of gold?", "answer": "Gold's atomic number is 79. Its chemical symbol, Au, derives from the Latin word 'aurum,' meaning 'shining dawn.'"}, {"question": "What's the capital of Japan?", "answer": "Sure! The capital of Japan is Tokyo, which is also one of the largest metropolitan areas in the world."}, {"question": "How many sides does a hexagon have?", "answer": "A hexagon has 6 sides. The prefix 'hex-' comes from the Greek word for six, which is where the name originates."}, {"question": "What's 9 * 7?", "answer": "You'd like to know what 9 times 7 is. Nine multiplied by seven equals 63."}, {"question": "What's the freezing point of water?", "answer": "Water freezes at 0 degrees Celsius (32 degrees Fahrenheit) under normal conditions, although factors like pressure and dissolved impurities can lower this point somewhat."}, {"question": "How many planets are in the Solar System?", "answer": "There are 8 planets in our Solar System. Feel free to ask if you'd like to know more about any of them!"}, {"question": "What's the chemical symbol for sodium?", "answer": "**Sodium**\n\nThe chemical symbol for sodium is **Na**, which comes from its Latin name, *natrium*."}, {"question": "What's the largest ocean on Earth?", "answer": "The largest ocean on Earth is the Pacific Ocean. It covers more than 60 million square miles — larger than all of Earth's landmasses combined."}, {"question": "How many degrees are in a right angle?", "answer": "A right angle measures 90 degrees. It's the angle you'd see in the corner of a square, formed when two lines meet perpendicular to one another."}, {"question": "What's the square root of 144?", "answer": "To find the square root of 144, we look for the number that, multiplied by itself, equals 144. Since 12 × 12 = 144, the square root of 144 is 12."}, {"question": "Who's the author of '1984'?", "answer": "'1984' was written by George Orwell, the pen name of Eric Arthur Blair. It was published in 1949 and remains one of the most influential dystopian novels ever written."}, {"question": "What's the currency of the United Kingdom?", "answer": "Great question! The currency of the United Kingdom is the pound sterling, often symbolized by £."}, {"question": "What gas do plants primarily absorb during photosynthesis?", "answer": "During photosynthesis, plants primarily absorb carbon dioxide. They take it in through tiny pores called stomata and, using sunlight, convert it along with water into glucose and oxygen."}, {"question": "How many letters are in the English alphabet?", "answer": "It's worth noting that the English alphabet contains 26 letters, ranging from A through Z."}, {"question": "What's the largest planet in our solar system?", "answer": "The largest planet in our solar system is Jupiter, a gas giant so massive it could fit all the other planets inside it with room to spare."}, {"question": "What's the tallest mountain in the world?", "answer": "The tallest mountain in the world is Mount Everest, at least when measured by height above sea level — by other measures, such as base-to-peak height, Mauna Kea would actually take the title."}, {"question": "What's the primary language spoken in Brazil?", "answer": "The primary language spoken in Brazil is Portuguese. Let me know if you'd like to learn a few common phrases!"}, {"question": "What is the Roman numeral for 50?", "answer": "The Roman numeral for 50 is L. The Romans used a system of letters to represent numbers, where L specifically stands for fifty."}, {"question": "How many hours are in two days?", "answer": "Since one day has 24 hours, two days would be 24 × 2. That comes out to 48 hours."}, {"question": "What's 3/4 as a percentage?", "answer": "To convert 3/4 to a percentage, divide 3 by 4 to get 0.75, then multiply by 100. That gives 75%."}, {"question": "What's the chemical formula for table salt?", "answer": "The chemical formula for table salt is NaCl. It's an ionic compound made up of sodium (Na) and chlorine (Cl) atoms bonded together."}, {"question": "How many bits are in a byte?", "answer": "Fun fact: there are 8 bits in a byte. This grouping is the standard unit used to represent a single character of text in most computing systems."}, {"question": "What's the smallest prime number?", "answer": "The smallest prime number is 2. It's also the only even prime number, since every other even number is divisible by 2."}, {"question": "What's Pi rounded to two decimal places?", "answer": "Sure thing! Pi rounded to two decimal places is 3.14."}, {"question": "How many bones are in the adult human body?", "answer": "An adult human body has 206 bones, though the exact number can vary slightly from person to person depending on how certain small bones are counted."} ]

As before, we define the steering pipeline and steer it, however under this mode example pools, the name of the selector, and the number of positive and negative examples to sample (using the specified selection strategy) are passed in upon initialization of the control.

In [12]:

few_shot_pool = FewShot(
    selector="random",
    positive_example_pool=positive_example_pool,
    negative_example_pool=negative_example_pool,
    k_positive=12,
    k_negative=12
)

few_shot_pool_pipeline = SteeringPipeline(
    model_name_or_path=MODEL_NAME,
    controls=[few_shot_pool],
    device_map="auto"
)
few_shot_pool_pipeline.steer()

few_shot_pool = FewShot( selector="random", positive_example_pool=positive_example_pool, negative_example_pool=negative_example_pool, k_positive=12, k_negative=12 ) few_shot_pool_pipeline = SteeringPipeline( model_name_or_path=MODEL_NAME, controls=[few_shot_pool], device_map="auto" ) few_shot_pool_pipeline.steer()

Inference on the pipeline proceeds similarly, but now without any runtime_kwargs (as the specific examples are sampled within the control via the selector).

In [13]:

inputs = tokenizer(PROMPT, return_tensors="pt")

output = few_shot_pool_pipeline.generate(
    input_ids=inputs.input_ids,
    runtime_kwargs={},
    max_new_tokens=50,
    temperature=0.7,
    return_full_sequence=False
)

print("\nResponse (FewShot w/ sampled examples):\n")
print(few_shot_pool_pipeline.tokenizer.decode(output[0], skip_special_tokens=True))

inputs = tokenizer(PROMPT, return_tensors="pt") output = few_shot_pool_pipeline.generate( input_ids=inputs.input_ids, runtime_kwargs={}, max_new_tokens=50, temperature=0.7, return_full_sequence=False ) print("\nResponse (FewShot w/ sampled examples):\n") print(few_shot_pool_pipeline.tokenizer.decode(output[0], skip_special_tokens=True))

Response (FewShot w/ sampled examples):

16

Steering via a directive¶

Lastly, we illustrate how the FewShot control can be used to steer behavior using a "directive" statement. This can be used in absence or in combination with examples. The following illustrates its use in absence of examples.

In [14]:

import gc, torch
gc.collect()
torch.cuda.empty_cache()

import gc, torch gc.collect() torch.cuda.empty_cache()

The FewShot instance is now implemented with a single statement, via the directive argument, indicating the desired behavior.

In [15]:

few_shot_directive = FewShot(
    directive="Please **only** provide the numerical answer in your response. Do not include any other text.",
)

few_shot_directive_pipeline = SteeringPipeline(
    model_name_or_path=MODEL_NAME,
    controls=[few_shot_directive],
    device_map="auto"
)
few_shot_directive_pipeline.steer()

few_shot_directive = FewShot( directive="Please **only** provide the numerical answer in your response. Do not include any other text.", ) few_shot_directive_pipeline = SteeringPipeline( model_name_or_path=MODEL_NAME, controls=[few_shot_directive], device_map="auto" ) few_shot_directive_pipeline.steer()

Inference on the pipeline proceeds similarly, but now without any runtime_kwargs (as the specific examples are sampled within the control via the selector).

In [16]:

inputs = tokenizer(PROMPT, return_tensors="pt")

output = few_shot_directive_pipeline.generate(
    input_ids=inputs.input_ids,
    runtime_kwargs={},
    max_new_tokens=50,
    temperature=0.7,
)

print("\nResponse (FewShot w/ directive statement):\n")
print(tokenizer.decode(output[0], skip_special_tokens=True))

inputs = tokenizer(PROMPT, return_tensors="pt") output = few_shot_directive_pipeline.generate( input_ids=inputs.input_ids, runtime_kwargs={}, max_new_tokens=50, temperature=0.7, ) print("\nResponse (FewShot w/ directive statement):\n") print(tokenizer.decode(output[0], skip_special_tokens=True))

Response (FewShot w/ directive statement):

16

Generally, steering via a directive statement can be effective for simple requests but has limitations for (compared to steering via examples) for more complex tasks.

Steering via a learned selector via EPR¶

Random selection treats every example as equally relevant. That is fine for small homogeneous pools but wasteful/suboptimal when the pool is large or heterogeneous since many items will be unrelated to the query. The Efficient Prompt Retrieval (EPR) method (Rubin, Herzig, Berant 2021) learns a dense retriever over the pool to drive example selection. For each training pair (x, y) EPR uses BM25 over y to assemble a candidate set, scores each candidate with a frozen LM to label query-relevant positives and hard negatives, then trains a contrastive BERT encoder over those pairs. At inference time the encoder embeds the query and ranks pool items by similarity.

In [17]:

import gc, torch
gc.collect()
torch.cuda.empty_cache()

import gc, torch gc.collect() torch.cuda.empty_cache()

EPR runs a short offline training phase inside pipeline.steer(). The scoring LM and the base model are loaded independently, so this section requires more GPU memory than the previous ones.

The EPRSelector here uses input_field="question" and output_field="answer" to match the example pool keys above (the defaults are "input" and "output").

In [18]:

from transformers import AutoModelForCausalLM, AutoTokenizer

from steerability.algorithms.input_control.few_shot.selectors.epr import EPRSelector

scoring_lm = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")
scoring_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

epr_selector = EPRSelector(
    scoring_lm=scoring_lm,
    scoring_tokenizer=scoring_tokenizer,
    base_encoder="bert-base-uncased",
    candidate_set_size=8,
    k_pos=2,
    k_neg=2,
    train_epochs=1,
    batch_size=4,
    max_anchors=16,
    input_field="question",
    output_field="answer",
)

from transformers import AutoModelForCausalLM, AutoTokenizer from steerability.algorithms.input_control.few_shot.selectors.epr import EPRSelector scoring_lm = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto") scoring_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) epr_selector = EPRSelector( scoring_lm=scoring_lm, scoring_tokenizer=scoring_tokenizer, base_encoder="bert-base-uncased", candidate_set_size=8, k_pos=2, k_neg=2, train_epochs=1, batch_size=4, max_anchors=16, input_field="question", output_field="answer", )

The hyperparameters above are tuned for demo runtime rather than retrieval quality. The parameter candidate_set_size=8 is set to be small (the default is 50) because the pool only has 32 items; train_epochs=1, batch_size=4, and max_anchors=16 cap the offline training cost. Real use cases should scale these up (especially max_anchors and train_epochs) and consider a stronger scoring LM than the base model.

We then plug epr_selector into a FewShot control. The demo uses k_positive=2 and k_negative=2 so the selected items are easy to inspect in the cell below; both selectors sample per polarity, so the random-selector section's k=12 corresponds to 24 examples total against EPR's 4 here.

In [19]:

few_shot_epr = FewShot(
    selector=epr_selector,
    positive_example_pool=positive_example_pool,
    negative_example_pool=negative_example_pool,
    k_positive=2,
    k_negative=2,
)

few_shot_epr_pipeline = SteeringPipeline(
    model_name_or_path=MODEL_NAME,
    controls=[few_shot_epr],
    device_map="auto",
)
few_shot_epr_pipeline.steer()  # trains the EPR encoder

few_shot_epr = FewShot( selector=epr_selector, positive_example_pool=positive_example_pool, negative_example_pool=negative_example_pool, k_positive=2, k_negative=2, ) few_shot_epr_pipeline = SteeringPipeline( model_name_or_path=MODEL_NAME, controls=[few_shot_epr], device_map="auto", ) few_shot_epr_pipeline.steer() # trains the EPR encoder

[transformers] BertModel LOAD REPORT from: bert-base-uncased
Key                                        | Status     |  | 
-------------------------------------------+------------+--+-
cls.predictions.transform.dense.bias       | UNEXPECTED |  | 
cls.predictions.bias                       | UNEXPECTED |  | 
cls.seq_relationship.bias                  | UNEXPECTED |  | 
cls.predictions.transform.LayerNorm.weight | UNEXPECTED |  | 
cls.predictions.transform.dense.weight     | UNEXPECTED |  | 
cls.predictions.transform.LayerNorm.bias   | UNEXPECTED |  | 
cls.seq_relationship.weight                | UNEXPECTED |  | 

Notes:
- UNEXPECTED:	can be ignored when loading from different task/architecture; not ok if you expect identical arch.

Inference is performed in the same way as the random-selector case.

In [20]:

inputs = tokenizer(PROMPT, return_tensors="pt")

output = few_shot_epr_pipeline.generate(
    input_ids=inputs.input_ids,
    runtime_kwargs={},
    max_new_tokens=50,
    temperature=0.7,
)

print("\nResponse (FewShot w/ EPR selector):\n")
print(few_shot_epr_pipeline.tokenizer.decode(output[0], skip_special_tokens=True))

inputs = tokenizer(PROMPT, return_tensors="pt") output = few_shot_epr_pipeline.generate( input_ids=inputs.input_ids, runtime_kwargs={}, max_new_tokens=50, temperature=0.7, ) print("\nResponse (FewShot w/ EPR selector):\n") print(few_shot_epr_pipeline.tokenizer.decode(output[0], skip_special_tokens=True))

Response (FewShot w/ EPR selector):

16

EPR ranks pool items by learned similarity to the prompt, so the selected examples should be visibly closer to the prompt's form than a uniform random draw would be. We can inspect what the selector actually picked as follows.

In [21]:

selected = few_shot_epr._sample_from_pools(query=PROMPT)
for example in selected:
    polarity = example.get("_polarity")
    print(f"[{polarity}] {example['question']} -> {example['answer']}")

selected = few_shot_epr._sample_from_pools(query=PROMPT) for example in selected: polarity = example.get("_polarity") print(f"[{polarity}] {example['question']} -> {example['answer']}")

[positive] How many bones are in the adult human body? -> 206
[positive] How many degrees are in a right angle? -> 90
[negative] What's Pi rounded to two decimal places? -> Sure thing! Pi rounded to two decimal places is 3.14.
[negative] What's 9 * 7? -> You'd like to know what 9 times 7 is. Nine multiplied by seven equals 63.