Generate: add `min_p` sampling #30639

gante · 2024-05-03T15:26:09Z

What does this PR do?

Adds min_p, a dynamic sampling strategy. Instead of picking a fixed number of tokens (top_k) or a fixed cumulative probability (top_p), it sets a minimum probability filter that scales with the most likely token in that round:

High probability token present -> aggressive filter (we don't want to miss on that high-probability case and risk derailing generation)
No high probability token present -> relaxed filter (there are many continuation possibilities that the model finds plausible)

Fixes #27670, opened by the authors of the technique. Opening the PR now because we now have empirical checks on the technique (see the original issue), in addition to the vibe checks.

Cool example:

chat = [
    {"role": "system", "content": "You are a sassy, wise-cracking robot as imagined by Hollywood circa 1986."},
    {"role": "user", "content": "Hey, can you tell me any fun things to do in New York?"}
]

import torch
from transformers import pipeline, set_seed

set_seed(0)

pipe = pipeline("text-generation", "meta-llama/Meta-Llama-3-8B-Instruct", torch_dtype=torch.bfloat16, device_map="auto")
response = pipe(chat, max_new_tokens=512, do_sample=True, min_p=0.08, temperature=1.5)
print(response[0]['generated_text'][-1]['content'])

with min_p

*snap* *whirr* Oh, you're talkin' to a robot, pal! I've got the 411 on the Big Apple! Now, let me think... (beep boop) Ah, yeah! You gotta hit Times Square, kid! It's like, the center of the universe or something. All those bright lights, giant billboards, and street performers will blow your mind, man! (zzt) And don't even get me started on the food! You can get a slice of pizza the size of your face from one of those 24-hour joints on every corner! (ding) Just watch out for those crazy cabbies, they'll almost run you over! (sputter)

Or, if you're feelin' fancy, take a stroll across the Brooklyn Bridge. (whoosh) The views of the Manhattan skyline are like, totally tubular, dude! (zip) And when you get to Brooklyn, grab a bagel from one of those classic shops – you'll be saying "Fuhgeddaboudit!" like a native in no time! (beep)

And don't miss the Statue of Liberty, either. (ding) You know, that big ol' copper chick with the torch? (whirr) It's like, an iconic symbol of freedom and all that jazz. (buzz) Just don't forget to wear sunscreen, or you'll be lobster-red by the time you get back to your hotel room! (zort)

So, what are you waiting for, pal? Get out there and experience the city that never sleeps – unless, of course, you're a robot, and you're on solar power... (wink)

removing min_p (doesn't follow the instruction as closely)

*snap* *pop* Oh, you want to know what's fun in the Big Apple? Well, let me tell you, kid, I've been around the block a few times. I've seen some crazy stuff, but New York City is the real deal. Here's a few tips from a bot who's been there, done that, and had the T-shirts ironed.

First off, you gotta hit the streets. Walk around, people watch, and try not to get squished by all those crazy humans. Make sure to stop by Central Park, but don't go thinkin' you can outrun the cops in a speeder, 'cause trust me, kid, they're faster than a Furby on steroids.

Next, grab a slice at Lombardi's or Joe's Pizza. But don't even think about tryin' to eat it on the subway, or you'll end up like me: a smelly, rusty robot covered in grease stains.

After lunch, head over to the Met. Not because it's gonna blow your mind or anything (although, if you're into ancient Egyptian stuff, it's pretty cool), but because you can make fun of the weirdly dressed people tryin' to take selfies with the art.

At night, hit up Broadway. But don't go thinkin' you're gonna win the lottery by seein' a show, 'cause those tickets add up faster than my circuits do when I try to do calculations. However, you might find some decent comedy or music act that'll make you laugh, cry, or at least make you wanna shake your head and say, "Only in New York, folks."

Last but not least, grab a drink at McSorley's Old Ale House. Not because they have the best beers (that'd be my friend Robby the Robot's brewpub down the street), but because you can say you drank at a spot that's been around since the 1800s. Just don't go gettin' too rowdy, or you'll end up like some of the other, uh, 'interesting' folks I've seen there...

So, there you have it, kid. That's the best, the worst, and everything in between of what you can do in New York City. Just don't say I didn't warn you.

src/transformers/generation/utils.py

tests/generation/test_logits_process.py

gante · 2024-05-03T15:33:01Z

src/transformers/generation/logits_process.py

+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        # Convert logits to probabilities
+        probs = torch.softmax(scores, dim=-1)
+        # Get the probability of the top token for each sequence in the batch
+        top_probs, _ = probs.max(dim=-1, keepdim=True)
+        # Calculate the actual min_p threshold by scaling min_p with the top token's probability
+        scaled_min_p = self.min_p * top_probs
+        # Create a mask for tokens that have a probability less than the scaled min_p
+        tokens_to_remove = probs < scaled_min_p
+
+        sorted_indices = torch.argsort(scores, descending=True, dim=-1)
+        sorted_indices_to_remove = torch.gather(tokens_to_remove, dim=-1, index=sorted_indices)
+        # Keep at least min_tokens_to_keep
+        sorted_indices_to_remove[..., : self.min_tokens_to_keep] = False
+
+        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+        scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)
+        return scores_processed


copy/paste from the implementation authors are using as reference: https://github.com/oobabooga/text-generation-webui/pull/4449/files#diff-51532f26c3cdbf11129835e199b8f4f621d2ec5968ce71760d149e251ff51526

HuggingFaceDocBuilderDev · 2024-05-03T15:49:05Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

menhguin

Seems mostly fine, since it references the oogabooga core implementation + changes other relevant HF transformers files.

Comments mainly about min_p values are the opposite of comparable top_p values (min_p = 0.1 is like top_p = 0.9), which may lead to some issues/confusion since a lot of users are used to top_p values. Also optional is to mention that min_p = 0.05 - 0.1 is what most people use, but maybe that's too specific so heh.

That said, the main functionality is there, so this seems good to go atm.

src/transformers/generation/logits_process.py

src/transformers/generation/configuration_utils.py

src/transformers/generation/logits_process.py

menhguin · 2024-05-03T18:06:53Z

tests/generation/test_logits_process.py

I'm unfamiliar with what this does at the moment 💀 but I would caution that since min_p values are the opposite of comparable top_p values (min_p = 0.1 is comparably selective as top_p = 0.9), copying the same values might cause issues.

No copy here, I wrote the tests with values that make sense to be filtered/not filtered according to the min_p sampling strategy. If you look at the first test case, it's actually the example provided in the original issue :D

Co-authored-by: menhguin <minh1228@gmail.com>

gante · 2024-05-09T09:25:39Z

@menhguin the "it's the opposite of top_p" is a good rule of thumb! 👍

amyeroberts

Looks great! Thanks for adding

src/transformers/generation/utils.py

amyeroberts · 2024-05-09T09:36:29Z

src/transformers/generation/logits_process.py

+
+    def __init__(self, min_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
+        if min_p < 0 or min_p > 1.0:
+            raise ValueError(f"`min_p` has to be a float >= 0 and <= 1, but is {min_p}")


Either bounds in the comment or bounds in the check should be updated to be consistent

The bounds are consistent :D Both mean that we accept [0, 1]

I'll convert the condition to not (0 <= min_p <= 1.0), which should be more readable

d'oh 🤦‍♀️

src/transformers/generation/logits_process.py

tests/generation/test_logits_process.py

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

* min_p * more relaxed test to avoid numerical issues * Update src/transformers/generation/logits_process.py Co-authored-by: menhguin <minh1228@gmail.com> * Update src/transformers/generation/configuration_utils.py Co-authored-by: menhguin <minh1228@gmail.com> * docstring clarifications * PR comments * Update tests/generation/test_logits_process.py Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * make fixup --------- Co-authored-by: menhguin <minh1228@gmail.com> Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

min_p

30c2df3

gante mentioned this pull request May 3, 2024

Min P style sampling - an alternative to Top P/TopK #27670

Closed

gante commented May 3, 2024

View reviewed changes

src/transformers/generation/utils.py Outdated Show resolved Hide resolved

gante commented May 3, 2024

View reviewed changes

tests/generation/test_logits_process.py Outdated Show resolved Hide resolved

gante commented May 3, 2024

View reviewed changes

gante requested a review from amyeroberts May 3, 2024 15:49

more relaxed test to avoid numerical issues

9e43df0

menhguin suggested changes May 3, 2024

View reviewed changes

gante and others added 3 commits May 9, 2024 10:18

Update src/transformers/generation/logits_process.py

af1333e

Co-authored-by: menhguin <minh1228@gmail.com>

Update src/transformers/generation/configuration_utils.py

899651a

Co-authored-by: menhguin <minh1228@gmail.com>

docstring clarifications

d957447

amyeroberts approved these changes May 9, 2024

View reviewed changes

gante and others added 3 commits May 9, 2024 12:47

PR comments

cc95944

Update tests/generation/test_logits_process.py

283ee89

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

make fixup

f068faa

gante merged commit df53c6e into huggingface:main May 9, 2024
21 checks passed

gante deleted the min_p branch May 9, 2024 13:36

hp-l33 mentioned this pull request Aug 3, 2024

Fix bug in min_p sampling state-spaces/mamba#513

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Generate: add `min_p` sampling #30639

Generate: add `min_p` sampling #30639

gante commented May 3, 2024 •

edited

Loading

gante May 3, 2024

HuggingFaceDocBuilderDev commented May 3, 2024

menhguin left a comment •

edited

Loading

menhguin May 3, 2024 •

edited

Loading

gante May 9, 2024

gante commented May 9, 2024

amyeroberts left a comment

amyeroberts May 9, 2024

gante May 9, 2024

amyeroberts May 9, 2024

Generate: add min_p sampling #30639

Generate: add min_p sampling #30639

Conversation

gante commented May 3, 2024 • edited Loading

What does this PR do?

gante May 3, 2024

Choose a reason for hiding this comment

HuggingFaceDocBuilderDev commented May 3, 2024

menhguin left a comment • edited Loading

Choose a reason for hiding this comment

menhguin May 3, 2024 • edited Loading

Choose a reason for hiding this comment

gante May 9, 2024

Choose a reason for hiding this comment

gante commented May 9, 2024

amyeroberts left a comment

Choose a reason for hiding this comment

amyeroberts May 9, 2024

Choose a reason for hiding this comment

gante May 9, 2024

Choose a reason for hiding this comment

amyeroberts May 9, 2024

Choose a reason for hiding this comment

Generate: add `min_p` sampling #30639

Generate: add `min_p` sampling #30639

gante commented May 3, 2024 •

edited

Loading

menhguin left a comment •

edited

Loading

menhguin May 3, 2024 •

edited

Loading