qwen2 models causing indexerror

[ml2s]

I'm trying to run the example quick_overview.ipynb on m1 mbp.

from dotenv import load_dotenv
from pathlib import Path
from byaldi import RAGMultiModalModel

load_dotenv()
# print(os.environ["HF_TOKEN"])

# Initialize RAGMultiModalModel
model = RAGMultiModalModel.from_pretrained("vidore/colqwen2-v1.0", device="mps")

# Test indexing
metadata = [{"filename":file_name} for file_name in os.listdir("docs")]

index_name = "attention_index"
model.index(
    input_path=Path("docs/"),
    index_name=index_name,
    store_collection_with_index=False,
    metadata=metadata,
    overwrite=True
)

# BLEU tables are on page 8 and 9. We've indexed the pdf and its evil mustached twin, so we should see similar scores occur twice for every relevant page.
query = "what's the BLEU score of this new strange method?"
results = model.search(query, k=5)

print(f"Search results for '{query}':")
for result in results:
    print(f"Doc ID: {result.doc_id}, Page: {result.page_num}, Score: {result.score}")

print("Test completed successfully!")

But got the following error:

---

IndexError Traceback (most recent call last)
Cell In[6], line 5
2 metadata = [{"filename":file_name} for file_name in os.listdir("docs")]
4 index_name = "attention_index"
----> 5 model.index(
6 input_path=Path("docs/"),
7 index_name=index_name,
8 store_collection_with_index=False,
9 metadata=metadata,
10 overwrite=True
11 )
13 # BLEU tables are on page 8 and 9. We've indexed the pdf and its evil mustached twin, so we should see similar scores occur twice for every relevant page.
14 query = "what's the BLEU score of this new strange method?"
... ...

File /opt/homebrew/Caskroom/miniforge/base/envs/colpali/lib/python3.12/site-packages/transformers/models/qwen2_vl/modeling_qwen2_vl.py:404, in VisionSdpaAttention.forward(self, hidden_states, cu_seqlens, rotary_pos_emb)
402 k = k.transpose(0, 1)
403 v = v.transpose(0, 1)
--> 404 attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0)
405 attn_output = attn_output.transpose(0, 1)
406 attn_output = attn_output.reshape(seq_length, -1)

IndexError: Dimension out of range (expected to be in range of [-3, 2], but got 3)

It seems the colpali models works fine, but qwen2 models cause the same error. Thanks for any help how to fix this.

[TonkaFails]

ColQwen2 on mps works with a different attention implementation. I don't know if you can change it using byaldi.

Example using colpali directly: Example

and use this:

ColQwen2ForRAG.from_pretrained(
model_name,
torch_dtype=torch.bfloat16,
device_map=device,
attn_implementation="eager",
low_cpu_mem_usage=True
),