custom stopping_critriea function doesn't receive logits scores (receives None instead)

Gandalf098 commented 1 year ago

System Info

transformers version: 4.29.2
Platform: Linux-5.15.107+-x86_64-with-glibc2.31
Python version: 3.10.11
Huggingface_hub version: 0.14.1
Safetensors version: not installed
PyTorch version (GPU?): 2.0.1+cu118 (True)
Tensorflow version (GPU?): 2.12.0 (True)
Flax version (CPU?/GPU?/TPU?): 0.6.9 (gpu)
Jax version: 0.4.8
JaxLib version: 0.4.7
Using GPU in script?: No
Using distributed or parallel set-up in script?: No

Who can help?

No response

Information

[ ] The official example scripts
[X] My own modified scripts

Tasks

[X] An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
[ ] My own task or dataset (give details below)

Reproduction

Reproduction Steps:

Initialize a BART model & its tokenizer (in my case it is facebook/bart-large)
Create a custom stopping_criteria function and add it to StoppingCriteriaList object
Run model.generate() with the your stopping criteria list as argument

Scores argument is always None

Example code:

import torch
from transformers import StoppingCriteriaList, BartForConditionalGeneration, BartTokenizer

def custom_stopping_criteria(input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
    print ("Scores:", scores)
    return False

stopping_criteria = StoppingCriteriaList([custom_stopping_criteria])

model = BartForConditionalGeneration.from_pretrained("facebook/bart-large", forced_bos_token_id=0)
tok = BartTokenizer.from_pretrained("facebook/bart-large")

example_english_phrase = "UN Chief Says There Is No <mask> in Syria"
batch = tok(example_english_phrase, return_tensors="pt")

model.generate(batch["input_ids"], stopping_criteria=stopping_criteria)

The above code uses a stopping critriea that just prints the scores value when called (which prints None)

Expected behavior

The expected behavior should be to have Scores logits populated with values instead of being None (values before or after softmax don't matter)

sgugger commented 1 year ago

cc @gante

gante commented 1 year ago

Hey @Gandalf098 (the white, I hope ;) )

By default, the scores are not initialized and are kept as None (see here). To enable score-keeping, you must pass return_dict_in_generate=True, output_scores=True to your .generate() call.

import torch
from transformers import StoppingCriteriaList, BartForConditionalGeneration, BartTokenizer

def custom_stopping_criteria(input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
    print("Scores:", scores)
    return False

stopping_criteria = StoppingCriteriaList([custom_stopping_criteria])

model = BartForConditionalGeneration.from_pretrained("facebook/bart-large", forced_bos_token_id=0)
tok = BartTokenizer.from_pretrained("facebook/bart-large")

example_english_phrase = "UN Chief Says There Is No <mask> in Syria"
batch = tok(example_english_phrase, return_tensors="pt")

model.generate(batch["input_ids"], stopping_criteria=stopping_criteria, return_dict_in_generate=True, output_scores=True)

ylacombe commented 1 year ago

Hi @gante and @Gandalf098,

According to the StoppingCriteria.__call__ signature and to its docstring, scores is supposed to be a torch.FloatTensor.

scores (torch.FloatTensor of shape (batch_size, config.vocab_size)) — Prediction scores of a language modeling head.

It makes sense to think of it as the last prediction scores of the language modeling head, meaning that the score-keeping here refers not to score (optional history of the prediction scores) but to next_token_scores (always available last prediction scores - at least for greedy decoding, we should verify for other decoding strategies).

In that sense, I do think we should correct this point. What do you think @gante?

gante commented 1 year ago

We might want to build some stopping criteria based on a sequence of tokens/sequence of scores, so this API is more general 🤗

We do need better docs and/or input validation, though, to detect these issues in advance. It is my priority for this month (and I'm keeping this issue open so I don't forget to address this case)

huggingface / transformers