Model doesn't generate output of prefix_allowed_tokens_fn.
Author: Chanwhistle
Pull request: no

### System Info

Hi, I use constrained decoding with prefix_allowed_tokens_fn.([GENRE github](https://github.com/facebookresearch/GENRE))

However, if prefix_allowed_tokens_fn function returns specific token like [1437],
model generate other tokens like [0] or [16] ...etc

For example
input_sent = [35524, 646]
trie.get(input_sent) = [1437] (constrained token by trie, that return to generate function)
generated_sent = [35524, 646, 8] or [35524, 646, 9] ...etc

I think generated_sent should be [35524, 646, 1437].

This only happens when I use bart like model, Bart, BioBART etc. Not in T5 or other else.
I tried num_beams, top_k or other methods but didn't work.
Could somebody knows why this happens and how can I solve this problem?
Really need help!

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

1. make trie file
2. use constrained decoding using Bart like models.
3. this happens when constrained decoding is more than 2 times in sentence.

example;

[2, 25522, 1308, 1242, 10003, 14554, 26764, 16628, _**35524, 646, 1437,**_ 50266, 27779, 36, 25522, 18695, _**35524, 646, 0**,_ 0, 0, 0, 0, 0, ....]

The first bold part doesn't have a bug, but something went wrong with the second part.

### Expected behavior

want to know reason why this happens

Comments:
Hi @Chanwhistle, thanks for opening this issue! 

In terms of reproduction, could you provide a minimal code snippet which reproduces the error? We get many issues and requests per day, so that we can address them in a timely manner, we need you to help us help you 🤗 

cc @gante 

Thank you!

code is like below.

```
        outputs = self.generate(
            **input_args,
            min_length=0,
            max_new_tokens=256,
            num_beams=num_beams,
            temperature=0.1,
            num_return_sequences=num_return_sequences,
            output_scores=False,
            return_dict_in_generate=True,
            forced_bos_token_id=None,
            prefix_allowed_tokens_fn=end_to_end_prefix_allowed_tokens_fn()
        )
```

input is 
"In a South African girl of Xhosa stock with severe { piebaldism } [ <D> ] and profound congenital { sensorineural deafness } [ <D> ]"

and expected output is

"In a South African girl of Xhosa stock with severe { piebaldism } [ <D> (words from consrtained decoding) ] and profound congenital { sensorineural deafness } [ <D> (words from consrtained decoding) ]"

But when I do this with my code, real output is 

"In a South African girl of Xhosa stock with severe { piebaldism } [ <D> piebaldism ] and profound congenital { sensorineural deafness } [ with with with with with .."

Constrained decoding doesn't done when the second constrained decoding happens. And also This happens only in Bart like model.


prefix_allowed_tokens function is like below.
Those code are from ([Facebook/Genre](https://github.com/facebookresearch/GENRE))

```
def end_to_end_prefix_allowed_tokens_fn(
    model,
    encode_fn,
    decode_fn,
    bos_token_id,
    pad_token_id,
    eos_token_id,
    vocabulary_length,
    task,
    type_dictionary,
    sentences: List[str],
    start_mention_token="{",
    end_mention_token="}",
    start_entity_token="[",
    end_entity_token="]",
    mention_trie: Trie = None,
    candidates_trie: Trie = None,
    mention_to_candidates_dict: Dict[str, List[str]] = None,
):

    assert not (
        candidates_trie is not None and mention_to_candidates_dict is not None
    ), "`candidates_trie` and `mention_to_candidates_dict` cannot be both != `None`"

    dict4ner = {}

    type_dictionary = {
            "disease_type_token" : "<D>",
        }
    
    model_type = str(model)
      codes = {
          n: encode_fn(" {}".format(c))[1]  
          for n, c in zip(
              (
                  "start_mention_token",
                  "end_mention_token",
                  "start_entity_token",
                  "end_entity_token",
              ),
              (
                  start_mention_token,
                  end_mention_token,
                  start_entity_token,
                  end_entity_token,
              ),
          )
      }
      codes["EOS"] = eos_token_id
      
      for key, value in type_dictionary.items():
          codes[key] = encode_fn(value)[1]
                   
      dict4ner = {
          v: encode_fn(" }} [ {} ]".format(v))[1:]  
          for k, v in type_dictionary.items()
      }
      
      sent_origs = [[codes["EOS"]] + encode_fn(sent)[1:] for sent in sentences]
        
    if mention_trie is None:
        mention_trie = DummyTrieMention(
            [
                i
                for i in range(vocabulary_length)
                if i not in (bos_token_id, pad_token_id,)
            ]
        )

    if candidates_trie is None and mention_to_candidates_dict is None:
        candidates_trie = DummyTrieEntity(
            [
                i
                for i in range(vocabulary_length)
                if i not in (bos_token_id, pad_token_id,)
            ],
            codes,
        )
       
    def prefix_allowed_tokens_fn_EL(batch_id, sent):

        sent = sent.tolist()
        status = get_status(sent)
        sent_orig = sent_origs[batch_id]
                
        if status == "o" or status == "m":
            trie_out = get_trie_outside_EL(sent, sent_orig)
        elif status == "e":
            trie_out = get_trie_entity(sent, sent_orig)
        else:
            raise RuntimeError
        
        return trie_out
    

    def get_status(sent):
        c = [
            codes[e]
            for e in (
                "start_mention_token",
                "end_mention_token",
                "start_entity_token",
                "end_entity_token",
            )
        ]
        status = sum(e in c for e in sent) % 4
        
        if status == 0:
            return "o"
        elif status == 1:
            return "m"
        else:
            return "e"
                        

    def get_trie_outside_EL(sent, sent_orig):
        pointer_end = get_pointer_end_EL(sent, sent_orig)

        if pointer_end:
            return [sent_orig[pointer_end]]
        else:
            return []


    def get_pointer_end(sent, sent_orig):
        i = 0
        j = 0
        
        while i < len(sent):
            if sent[i] == sent_orig[j]:
                i += 1
                j += 1
            elif (
                sent[i] == codes["start_mention_token"] or 
                sent[i] == codes["end_mention_token"] 
            ):
                i += 1
            elif sent[i] == codes["start_entity_token"]:
                i += 1
                while sent[i] != codes["end_entity_token"]:
                    i += 1
                i += 1
            else:
                return None
            
        return j if j != len(sent_orig) else None
    
    def get_pointer_end_EL(sent, sent_orig):
        i = 0
        j = 0
        
        while i < len(sent):
            if sent[i] == sent_orig[j]:
                i += 1
                j += 1
            elif sent[i] != codes["end_entity_token"]:
                i += 1
            else:
                return None
            
        return j if j != len(sent_orig) else None
    
        
    def get_trie_mention(sent, sent_orig):
                
        pointer_start, _ = get_pointer_mention(sent)
        if pointer_start + 1 < len(sent):
            ment_next = mention_trie.get(sent[pointer_start + 1 :])
        else:
            ment_next = mention_trie.get([])

        pointer_end = get_pointer_end(sent, sent_orig)

        if pointer_end:
            if sent_orig[pointer_end] != codes["EOS"]:
                if sent_orig[pointer_end] in ment_next:
                    if codes["EOS"] in ment_next:
                        return [sent_orig[pointer_end], codes["end_mention_token"]]
                    else:
                        return [sent_orig[pointer_end]]
                elif codes["EOS"] in ment_next:
                    return [codes["end_mention_token"]]
                else:
                    return []
            else:
                return [codes["end_mention_token"]]
        else:
            return []
        
    def get_pointer_mention(sent):
        pointer_end = -1
        pointer_start = None
        
        for i, e in enumerate(sent):
            if e == codes["start_mention_token"]:
                pointer_start = i
            elif e == codes["end_mention_token"]:
                pointer_end = i
        # for debug
        if pointer_start == None:
            import pdb;pdb.set_trace()
            
        return pointer_start, pointer_end

    def get_trie_entity(sent, sent_orig):
        pointer_start, pointer_end = get_pointer_mention(sent)

        if pointer_start + 1 != pointer_end:
            mention = decode_fn(sent[pointer_start + 1 : pointer_end]).strip()

            if task == "NER":
                candidates_trie_tmp = Trie(dict4ner.values())
            elif candidates_trie is not None:
                candidates_trie_tmp = candidates_trie
            elif mention_to_candidates_dict is not None:
                candidates_trie_tmp = Trie(
                    [
                        encode_fn(
                            " {} {} {} {}".format(
                                end_mention_token,
                                start_entity_token,
                                e,
                                end_entity_token,
                            )
                        )[1:]
                        for e in mention_to_candidates_dict.get(mention, ["CUI-less"])
                    ]
                )
            else:
                raise RuntimeError()
            
            return candidates_trie_tmp.get(sent[pointer_end:])
        
        return []
 
    return prefix_allowed_tokens_fn_EL
```

Thanks!

Hi @Chanwhistle . I had similar problems with `prefix_allowed_tokens_fn` in the past. It's hard to say from your code what went wrong, but given that generation fails only in BART-like models, I would advise to check the tokenization. 

You can set a breakpoint to verify if the trie can correctly get the `allowed_tokens` given the generated sentence. Also, ensure that the special start and end of mention/entity tokens are tokenized separately, without being merged with adjacent tokens.

Hi @Chanwhistle 👋 Having a popular project like `transformers` means we get many support and feature requests — if we want to maximize how much we help the community, the community has to help us stay productive 🙏

To that end, please share a *short stand-alone script* where the issue is clearly reproducible on *any* computer. Thank you 🤗

(If this is your first issue with us, check [this guide](https://huggingface.co/course/chapter8/5?fw=pt).)

@gante Thank you for answer.

I read given webpage and made short stand-alone script, and checked same error occured.

- `transformers` version: 4.35.2
- Platform: Linux-5.4.0-169-generic-x86_64-with-glibc2.31
- Python version: 3.9.18
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.1
- Accelerate version: 0.25.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No



Here is the code.


```
from typing import Dict, List
import torch
from transformers import (
    BartTokenizer,
    BartForConditionalGeneration,
)

class Trie(object):
    def __init__(self, sequences: List[List[int]] = []):
        self.trie_dict = {}
        self.len = 0
        if sequences:
            for sequence in sequences:
                Trie._add_to_trie(sequence, self.trie_dict)
                self.len += 1

        self.append_trie = None
        self.bos_token_id = None

    def get(self, prefix_sequence: List[int]):
        return Trie._get_from_trie(
            prefix_sequence, self.trie_dict, self.append_trie, self.bos_token_id
        )

    @staticmethod
    def _add_to_trie(sequence: List[int], trie_dict: Dict):
        if sequence:
            if sequence[0] not in trie_dict:
                trie_dict[sequence[0]] = {}
            Trie._add_to_trie(sequence[1:], trie_dict[sequence[0]])

    @staticmethod
    def _get_from_trie(
        prefix_sequence: List[int],
        trie_dict: Dict,
        append_trie=None,
        bos_token_id: int = None,
    ):
        if len(prefix_sequence) == 0:
            output = list(trie_dict.keys())
            if append_trie and bos_token_id in output:
                output.remove(bos_token_id)
                output += list(append_trie.trie_dict.keys())
            return output
        elif prefix_sequence[0] in trie_dict:
            return Trie._get_from_trie(
                prefix_sequence[1:],
                trie_dict[prefix_sequence[0]],
                append_trie,
                bos_token_id,
            )
        else:
            if append_trie:
                return append_trie.get(prefix_sequence)
            else:
                return []

    def __len__(self):
        return self.len

    def __getitem__(self, value):
        return self.get(value)
    
def get_end_to_end_prefix_allowed_tokens_fn_hf(
    model,
    sentences: List[str],
    start_mention_token="{",
    end_mention_token="}",
    start_entity_token="[",
    end_entity_token="]",
    mention_trie: Trie = None,
    candidates_trie: Trie = None,
    mention_to_candidates_dict: Dict[str, List[str]] = None,
    task = None
):
    return _get_end_to_end_prefix_allowed_tokens_fn(
        model,
        lambda x: model.tokenizer.encode(x),
        lambda x: model.tokenizer.decode(torch.tensor(x)),
        model.tokenizer.bos_token_id,
        model.tokenizer.pad_token_id,
        model.tokenizer.eos_token_id,
        len(model.tokenizer) - 1,
        task,
        sentences,
        start_mention_token,
        end_mention_token,
        start_entity_token,
        end_entity_token,
        mention_trie,
        candidates_trie,
        mention_to_candidates_dict,
    )


def _get_end_to_end_prefix_allowed_tokens_fn(
    model,
    encode_fn,
    decode_fn,
    bos_token_id,
    pad_token_id,
    eos_token_id,
    vocabulary_length,
    task,
    sentences: List[str],
    start_mention_token="{",
    end_mention_token="}",
    start_entity_token="[",
    end_entity_token="]",
    mention_trie: Trie = None,
    candidates_trie: Trie = None,
    mention_to_candidates_dict: Dict[str, List[str]] = None,
):
    
    codes = {
        n: encode_fn(" {}".format(c))[1]  
        for n, c in zip(
            (
                "start_mention_token",
                "end_mention_token",
                "start_entity_token",
                "end_entity_token",
            ),
            (
                start_mention_token,
                end_mention_token,
                start_entity_token,
                end_entity_token,
            ),
        )
    }
    codes["EOS"] = eos_token_id
    codes["disease_type_token"] = encode_fn("<D>")[1]
              
    sent_origs = [[codes["EOS"]] + encode_fn(sent)[1:] for sent in sentences]

    def prefix_allowed_tokens_fn_EL(batch_id, sent):

        sent = sent.tolist()
        status = get_status(sent)
        sent_orig = sent_origs[batch_id]
                
        if status == "o" or status == "m":
            trie_out = get_trie_outside_EL(sent, sent_orig)
        elif status == "e":
            trie_out = get_trie_entity(sent, sent_orig)
            if trie_out == codes["EOS"]:
                trie_out = get_trie_outside_EL(sent, sent_orig)
        else:
            raise RuntimeError
        
        return trie_out
    

    def get_status(sent):
        c = [
            codes[e]
            for e in (
                "start_mention_token",
                "end_mention_token",
                "start_entity_token",
                "end_entity_token",
            )
        ]
        status = sum(e in c for e in sent) % 4
        
        if status == 0:
            return "o"
        elif status == 1:
            return "m"
        else:
            return "e"
                      
    def get_trie_outside_EL(sent, sent_orig):
        pointer_end = get_pointer_end_EL(sent, sent_orig)

        if pointer_end:
            return [sent_orig[pointer_end]]
        else:
            return []
    
    def get_pointer_end_EL(sent, sent_orig):
        i = 0
        j = 0
        
        while i < len(sent):
            if sent[i] == sent_orig[j]:
                i += 1
                j += 1
            elif sent[i] != codes["end_entity_token"]:
                i += 1
            else:
                return None
            
        return j if j != len(sent_orig) else None
            
    def get_pointer_mention(sent):
        pointer_end = -1
        
        for i, e in enumerate(sent):
            if e == codes["start_mention_token"]:
                pointer_start = i
            elif e == codes["end_mention_token"]:
                pointer_end = i
            
        return pointer_start, pointer_end

    def get_trie_entity(sent, sent_orig):
        pointer_start, pointer_end = get_pointer_mention(sent)

        if pointer_start + 1 != pointer_end:
            mention = decode_fn(sent[pointer_start + 1 : pointer_end]).strip()

            if candidates_trie is not None:
                candidates_trie_tmp = candidates_trie
            elif mention_to_candidates_dict is not None:
                candidates_trie_tmp = Trie(
                    [
                        encode_fn(
                            " {} {} {} {}".format(
                                end_mention_token,
                                start_entity_token,
                                e,
                                end_entity_token,
                            )
                        )[1:]
                        for e in mention_to_candidates_dict.get(mention, ["CUI-less"])
                    ]
                )
            else:
                raise RuntimeError()
            
            return candidates_trie_tmp.get(sent[pointer_end:])
        
        return []
    
    return prefix_allowed_tokens_fn_EL
    

class _GENREHubInterface:
    def sample(
        self,
        sentences: List[str],
        num_beams: int = 5,
        num_return_sequences=1,
        text_to_id: Dict[str, str] = None,
        marginalize: bool = False,
        **kwargs
    ) -> List[str]:
        
        input_args = {
            k: v.to(self.device) if isinstance(v, torch.Tensor) else v
            for k, v in self.tokenizer.batch_encode_plus(
                sentences, padding=True, return_tensors="pt"
            ).items()
        }
        
        outputs = self.generate(
            **input_args,
            min_length=0,
            max_new_tokens=128,
            num_beams=num_beams,
            num_return_sequences=num_return_sequences,
            early_stopping=False,
            output_scores=False,
            return_dict_in_generate=False,
            forced_bos_token_id=None,
            **kwargs
        )
        
        return outputs

    def encode(self, sentence):
        return self.tokenizer.encode(sentence, return_tensors="pt")[0]

class GENREHubInterface(_GENREHubInterface, BartForConditionalGeneration):
    pass
    
    
class GENRE(BartForConditionalGeneration):
    @classmethod
    def from_pretrained(cls, model_name_or_path):
        model = GENREHubInterface.from_pretrained(model_name_or_path)
        tokenizer = BartTokenizer.from_pretrained(model_name_or_path)
        start_disease_token = "<D>"
        tokenizer.add_tokens([start_disease_token])
        model.resize_token_embeddings(len(tokenizer))
        model.tokenizer = tokenizer
        
        return model
```

and Here is the script.

```
sentence = ["Germ-line and somatic truncating mutations of the { APC } [ <D> ] gene are thought to initiate { colorectal tumor } [ <D> ] formation in { familial adenomatous polyposis syndrome } [ <D> ] and sporadic colorectal carcinogenesis, respectively."]

model = GENRE.from_pretrained("facebook/bart-base")

prefix_allowed_tokens_fn = get_end_to_end_prefix_allowed_tokens_fn_hf(
    model,
    sentence,
    candidates_trie=Trie([
        model.encode(" }} [ <D> {} ]".format(e))[1:].tolist()
        for e in ["APC", "colorectal tumor", "familial adenomatous polyposis syndrome"]
    ])
)

output = model.sample(
    sentence,
    prefix_allowed_tokens_fn=prefix_allowed_tokens_fn ,
)

print(model.tokenizer.batch_decode(output, skip_special_tokens=True))

```
   
My expected output is 

["Germ-line and somatic truncating mutations of the { APC } [ <D> APC  ] gene are thought to initiate { colorectal tumor } [ <D> colorectal tumor ] formation in { familial adenomatous polyposis syndrome } [ <D> familial adenomatous polyposis syndrome ] and sporadic colorectal carcinogenesis, respectively."]

but this script's output is 

['Germ-line and somatic truncating mutations of the { APC } [  <D>  familial adenomatous polyposis syndrome ] gene are thought to initiate { colorectal tumor } [ and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and the']


Yeah no, that's not a short script 😅 I'm sorry, we will only dedicate resources with a short reproducer of the issue, without significant external code.



This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

@Chanwhistle I'm sorry to bother you, but I'm having the same problem with the code as you are. But I used the t5 model instead of bart. Have you solved the problem now? How is it solved? Thank you very much!


Hi @ERI2333 ,
Do you mean you have same problem in T5 model?

Yes @Chanwhistle , I have the same problem in T5 model. Is your problem still there? After that I found a related post , I'm not sure if it works and am trying it. I really reallyhope it works. https://github.com/facebookresearch/GENRE/issues/103

@ERI2333 Actually, In my case, this problem never appeared in my code with T5 model. 
Only Bart based model have this problem(Bart, BioBART etc..)
That makes me so confusing cuz I have any clue that why this problem happens.
And I checked my transformers version is v4.35.2.