Google RecurrentGemma Models don't work in Transformers 4.43 anymore

[dipanjanS]

System Info

Transformers version 4.43.3 causes this error, works fine in 4.42

from transformers import AutoTokenizer, AutoModelForCausalLM
import transformers
import torch

model_id = "google/recurrentgemma-2b-it"
dtype = torch.bfloat16
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="cuda",
    torch_dtype=dtype,
)

chat = [
    { "role": "user", "content": "Explain what is AI in 3 bullet points" },
]
prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)

inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
outputs = model.generate(input_ids=inputs.to(model.device),
                         max_new_tokens=150)
print(tokenizer.decode(outputs[0]))

Gives the following error now:

RecurrentGemmaForCausalLM.forward() got an unexpected keyword argument 'position_ids'

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

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

Reproduction

Already mentioned above.

Expected behavior

The prompt should run without errors just like in version 4.42

[ArthurZucker]

Opening a PR for a fix!

[ArthurZucker]

It was broken by #31549 😅

[dipanjanS]

Thanks for the prompt fix, much appreciated!