error in LangBridgeModel._reorder_cache function when beam search + use_cache=True

[jjkim0807]

Traceback (most recent call last):
File "text2sql_few_shot.py", line 248, in
generated_sqls = text2sql_func(model, input_seq, tokenizer, max_tokens, max_new_tokens, new_eos_token_id, enc_tokenizer=enc_tokenizer) File "text2sql_few_shot.py", line 159, in text2sql_func
generate_ids = model.generate(
File "/root/miniconda/envs/CodeS/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/root/miniconda/envs/CodeS/lib/python3.8/site-packages/transformers/generation/utils.py", line 1655, in generate
result = self._beam_search(
File "/root/miniconda/envs/CodeS/lib/python3.8/site-packages/transformers/generation/utils.py", line 3248, in _beam_search
model_kwargs["past_key_values"] = self._temporary_reorder_cache(
File "/root/miniconda/envs/CodeS/lib/python3.8/site-packages/transformers/generation/utils.py", line 2888, in _temporary_reorder_cache
past_key_values = self._reorder_cache(past_key_values, beam_idx)
File "/root/CodeS/LangBridge/langbridge/modeling_langbridge.py", line 421, in _reorder_cache
lm_past_beam = tuple(
TypeError: 'int' object is not iterable

When I try beam search with use_cache=True option, above error occurs.

generate_ids = model.generate( **inputs, max_new_tokens = max_new_tokens, num_beams = 4, num_return_sequences = 4, use_cache = True, eos_token_id = eos_token_id )

[MattYoon]

Hey @jjkim0807, thank you for your interest in our work.

We didn't use beam search in our experiments so the current code doesn't support beam search.

To use beam search you would need to rewrite the _reorder_cache method.

If you manage to get it to work, please open a PR and we are more than happy to incorporate your contribution!