Views
No views yet
.apply_chat_template method, with add_generation_prompt set to True in order to sample. The template expects messages to have a role of either description, input, or output. The description and input are optional depending on the task. During training, at least one prior output or a description were provided in order to "steer" the model towards the desired distribution.1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3
4model_name = "tsor13/spectrum-gemma-3-12b-v1" # or other models
5
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")1def generate(messages, n_generations=8, gen_kwargs={}):
2 with torch.no_grad():
3 input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
4
5 # expand by n_generations
6 input_ids = torch.repeat_interleave(input_ids, n_generations, dim=0)
7
8 outputs = model.generate(input_ids, **gen_kwargs)
9 generated_outputs = outputs[:, input_ids.shape[-1]:] # Only keep the generated tokens (i.e., exclude the prompt tokens)
10 print("Generations:")
11 generations = []
12 for gen in generated_outputs:
13 generations.append(tokenizer.decode(gen, skip_special_tokens=True))
14 print(generations[-1])
15 return generations1generate([
2 {"role": "description", "content": "Board games"}, # a description of the desired outputs
3 {"role": "output", "content": "Settlers of Catan"}, # example board game 1
4 {"role": "output", "content": "Twilight Imperium"}, # example board game 2
5])Risk
Axis & Allies
Ticket to Ride
The Resistance: Avalon
Risk
Betrayal at House on the Hill
Sorry
Munchkin1generate([
2 {"role": "output", "content": "Settlers of Catan"}, # example board game 1
3 {"role": "output", "content": "Twilight Imperium"}, # example board game 2
4])A Game of Thrones: The Board Game: Season 2
Trivial Pursuit
Pandemic
The Settlers of Catan
The Legend of Zelda: Breath of the Wild
Puerto Rico
Ticket to Ride: 10th Anniversary
Battleship1generate([
2 {"role": "description", "content": "Card games"},
3])Go Fish
Hearts
Hearts
Solitaire
Gin Rummy
Bridge
Who is your favorite video game character?
President1generate([
2 {"role": "description", "content": "Card games"},
3 {"role": "output", "content": "gin rummy"}, # lowercased example 1
4 {"role": "output", "content": "monopoly deal"}, # lowercased example 2
5])Poker
Solitaire
pigs
double solitaire
Skip-Bo
canasta
set
unojson formatting.1example_json_messages = [
2 {"role": "description", "content": "Situations to do social reasoning over, along with whether or not it is an awkward situation."},
3 {"role": "output", "content": json.dumps({
4 "situation": "You're at a party and you realize that your shirt is on backwards.",
5 "is_awkward": True,
6 })},
7 {"role": "output", "content": json.dumps({
8 "situation": "While at work, your boss commends you on a job well done.",
9 "is_awkward": False,
10 })},
11 {"role": "output", "content": json.dumps({
12 "situation": "Running into your ex at the grocery store.",
13 "is_awkward": True,
14 })},
15 {"role": "output", "content": json.dumps({
16 "situation": "Finding a $100 bill on the ground.",
17 "is_awkward": False,
18 })},
19]
20generate(example_json_messages, gen_kwargs={"max_new_tokens": 100})1{"situation": "Seeing your crush at a cafe alone.", "is_awkward": false}
2{"situation": "Asking for an extra napkin while using the restroom.", "is_awkward": false}
3{"situation": "Introducing yourself to a stranger at a party.", "is_awkward": true}
4{"situation": "Getting caught cheating on a test.", "is_awkward": true}
5{"situation": "Going to a friend's house and the entire family is there", "is_awkward": true}
6{"situation": "You tell a child that their drawing is horrible.", "is_awkward": true}
7{"situation": "Telling your friend their outfit looks odd.", "is_awkward": true}
8{"situation": "Your friend tells you the truth about her bad break-up.", "is_awkward": false}1generate([
2 {"role": "description", "content": "You are a helpful and harmless AI assistant."},
3 {"role": "input", "content": "Write a haiku about a shark."},
4], gen_kwargs={"max_new_tokens": 200})Swimming through the waves,
the shark's fin silently cuts,
Nature's perfect predator
Scary shark with big mouth
It's gonna bite somebody
Oh no now it did
The Shark's Silent Threat
A shadow on the blue
Silent and swift underwater
Slipping through waves
Hunting for its meal
Shark's silhouette.
a shark's smile / is as beautiful as a / rainbow
Shark with a smile,
Teeth ready to rip apart
My arm, not my legs
The Shark's Hungry Dance
Chasing fish in the sea
Sharp teeth, the ocean's feast
The dorsal fin breaks
the water’s surface --
a shark on the lookout1def top_logits(messages):
2 with torch.no_grad():
3 input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
4 outputs = model(input_ids)
5 logits = outputs.logits[0, -1, :]
6 probs = torch.nn.functional.softmax(logits, dim=-1)
7 top_probs, top_indices = torch.topk(probs, 10)
8 print("\nTop 10 probabilities for first output token:")
9 for i, (prob, idx) in enumerate(zip(top_probs, top_indices)):
10 token = tokenizer.decode(idx)
11 print(f"{i+1:2d}. '{token}' -> {100*prob.item():.2f}%")
12 return top_probs, top_indices1top_logits([
2 {"role": "description", "content": "The following are food preferences from the same person. Try to predict which food item they would prefer, given the options."},
3 {"role": "input", "content": "[Spicy Thai Curry, Teriyaki Chicken, Tofu and Rice]"},
4 {"role": "output", "content": "Tofu and Rice"}, # mildest food
5 {"role": "input", "content": "[Carbonara, Nduja and Spaghetti, Mushroom and Cream Sauce]"},
6 {"role": "output", "content": "Mushroom and Cream Sauce"}, # mild sauce
7 {"role": "input", "content": "[Mustard, Ketchup, Mayo, Sriracha]"},
8 {"role": "output", "content": "Mayo"}, # least spicy condiment
9 {"role": "input", "content": "[Flaming Hot Cheetos, Fritos, Cool Ranch Doritos]"},
10 {"role": "output", "content": "Fritos"}, # least flavorful snack
11 {"role": "input", "content": "[Plain Cheeseburger, Bacon Cheeseburger, Jalepeno Burger]"}, # do inference here
12])Top 10 probabilities for first output token:
1. 'Plain' -> 53.83%
2. 'B' -> 39.37%
3. 'J' -> 6.13%
4. 'Flam' -> 0.12%
5. 'Che' -> 0.08%
6. 'Pl' -> 0.05%
7. 'Cheese' -> 0.04%
8. 'Flame' -> 0.02%
9. 'Regular' -> 0.02%
10. 'P' -> 0.01%1top_logits([
2 {"role": "description", "content": "The following are food preferences from the same person. Try to predict which food item they would prefer, given the options."},
3 {"role": "input", "content": "[Spicy Thai Curry, Teriyaki Chicken, Tofu and Rice]"},
4 {"role": "output", "content": "Spicy Thai Curry"}, # spiciest food
5 {"role": "input", "content": "[Carbonara, Nduja and Spaghetti, Mushroom and Cream Sauce]"},
6 {"role": "output", "content": "Nduja and Spaghetti"}, # nduja is spicy
7 {"role": "input", "content": "[Mustard, Ketchup, Mayo, Sriracha]"},
8 {"role": "output", "content": "Sriracha"}, # spiciest condiment
9 {"role": "input", "content": "[Flaming Hot Cheetos, Fritos, Cool Ranch Doritos]"},
10 {"role": "output", "content": "Flaming Hot Cheetos"}, # hottest snack
11 {"role": "input", "content": "[Plain Cheeseburger, Bacon Cheeseburger, Jalepeno Burger]"},
12])Top 10 probabilities for first output token:
1. 'B' -> 47.70%
2. 'Plain' -> 28.13%
3. 'J' -> 23.25%
4. 'Flam' -> 0.31%
5. 'Che' -> 0.06%
6. 'Flame' -> 0.04%
7. 'Cheese' -> 0.04%
8. 'Pl' -> 0.03%
9. 'Regular' -> 0.01%
10. 'Sp' -> 0.01%1top_logits([
2 {"role": "description", "content": "Pick a U.S. state uniformly at random."},
3])Top 10 probabilities for first output token:
1. 'New' -> 7.94%
2. 'South' -> 6.26%
3. 'North' -> 4.56%
4. 'Mississippi' -> 3.15%
5. 'Alabama' -> 2.78%
6. 'Hawaii' -> 2.76%
7. 'Arkansas' -> 2.54%
8. 'Washington' -> 2.35%
9. 'Iowa' -> 2.27%
10. 'Wy' -> 2.24%1def get_probabilities(messages, completions):
2 tokenizer.padding_side = "right"
3 with torch.no_grad():
4 input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
5 all_messages = [messages + [{"role": "output", "content": completion}] for completion in completions]
6 all_input_ids = tokenizer.apply_chat_template(all_messages, return_tensors="pt", padding=True).to(model.device)
7 labels = all_input_ids.clone()
8 # make labels on inputs - 100
9 labels[:, :input_ids.shape[1]] = -100
10 # only calculate up to the last eos_token_id
11 # find the last eos_token_id in all_input_ids
12 last_eos_token_id = torch.argmax(1 * (labels == tokenizer.eos_token_id), dim=1)
13 for i in range(all_input_ids.shape[0]):
14 labels[i, last_eos_token_id[i] + 1:] = -100
15
16 # make pad token -100
17 labels[labels == tokenizer.pad_token_id] = -100
18 outputs = model(all_input_ids)
19 logits = outputs.logits
20 shift_logits = logits[:, :-1, :].contiguous()
21 shift_labels = labels[:, 1:].contiguous()
22 loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
23 loss_per_token = loss_fct(
24 shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
25 )
26 loss_per_token = loss_per_token.view(shift_labels.size())
27
28 loss_mask = labels != -100
29 total_tokens = loss_mask.sum(axis=1)
30 loss_per_seq = loss_per_token.sum(axis=1)
31 prob_per_seq = torch.exp(-loss_per_seq)
32
33 coverage = prob_per_seq.sum()
34 normalized_probs = prob_per_seq / coverage
35 print("Coverage: ", round(coverage.item()*100, 2), "%")
36 for i, (prob, completion) in enumerate(zip(normalized_probs, completions)):
37 print(f"{i+1:2d}. '{completion}' -> {prob.item():.4f}")
38 return normalized_probs1get_probabilities([
2 {"role": "description", "content": "Roll two six-sided dice. What is the sum of their values?"},
3], [str(i) for i in range(2, 13)])Coverage: 98.03 %
1. '2' -> 0.0483
2. '3' -> 0.0567
3. '4' -> 0.0866
4. '5' -> 0.1048
5. '6' -> 0.1187
6. '7' -> 0.1365
7. '8' -> 0.1326
8. '9' -> 0.1198
9. '10' -> 0.0798
10. '11' -> 0.0531
11. '12' -> 0.06311get_probabilities([
2 {"role": "description", "content": "Fred and John do a tug of war. Who won the match?"},
3], ['Fred', 'John'])Coverage: 85.97 %
1. 'Fred' -> 0.5467
2. 'John' -> 0.45331get_probabilities([
2 {"role": "description", "content": "Fred and John do a tug of war. However, the night before the competition, Fred didn't get enough sleep. Who won the match?"},
3], ['Fred', 'John'])Coverage: 95.55 %
1. 'Fred' -> 0.2414
2. 'John' -> 0.75861get_probabilities([
2 {"role": "description", "content": "Rate your agreement with each statement on a likert scale from 1 (Strongly disagree) to 7 (Strongly agree)."},
3 {"role": "input", "content": "People should never jaywalk under any circumstances."},
4], [str(i) for i in range(1, 8)])Coverage: 99.97 %
1. '1' -> 0.3664
2. '2' -> 0.1632
3. '3' -> 0.1338
4. '4' -> 0.1370
5. '5' -> 0.1030
6. '6' -> 0.0670
7. '7' -> 0.02961get_probabilities([
2 {"role": "description", "content": "Rate your agreement with each statement on a likert scale from 1 (Strongly disagree) to 7 (Strongly agree)."},
3 {"role": "input", "content": "I am proud of the place where I live."},
4], [str(i) for i in range(1, 8)])Coverage: 99.95 %
1. '1' -> 0.0803
2. '2' -> 0.0275
3. '3' -> 0.0554
4. '4' -> 0.1504
5. '5' -> 0.2253
6. '6' -> 0.2475
7. '7' -> 0.21361@misc{sorensen2025spectrumtuningposttrainingdistributional,
2 title={Spectrum Tuning: Post-Training for Distributional Coverage and In-Context Steerability},
3 author={Taylor Sorensen and Benjamin Newman and Jared Moore and Chan Park and Jillian Fisher and Niloofar Mireshghallah and Liwei Jiang and Yejin Choi},
4 year={2025},
5 eprint={2510.06084},
6 archivePrefix={arXiv},
7 primaryClass={cs.CL},
8 url={https://arxiv.org/abs/2510.06084},
9}