fix AttributeError: 'NoneType' object has no attribute 'dim'

[konrad-woj]

Description

Running lightrag_hf_demo with nomic-embed-text raises AttributeError: 'NoneType' object has no attribute 'dim'. version: 1.2.1

File "/Users/konradwojciechowski/PyCharmProjects/graph-rag/LightRAG/lightrag/lightrag.py", line 982, in query
    return loop.run_until_complete(self.aquery(query, param, system_prompt))  # type: ignore
  File "/Users/konradwojciechowski/.pyenv/versions/3.10.16/lib/python3.10/asyncio/base_events.py", line 649, in run_until_complete
    return future.result()
  File "/Users/konradwojciechowski/PyCharmProjects/graph-rag/LightRAG/lightrag/lightrag.py", line 1023, in aquery
    response = await naive_query(
  File "/Users/konradwojciechowski/PyCharmProjects/graph-rag/LightRAG/lightrag/operate.py", line 1564, in naive_query
    results = await chunks_vdb.query(query, top_k=query_param.top_k)
  File "/Users/konradwojciechowski/PyCharmProjects/graph-rag/LightRAG/lightrag/kg/nano_vector_db_impl.py", line 83, in query
    embedding = await self.embedding_func([query])
  File "/Users/konradwojciechowski/PyCharmProjects/graph-rag/LightRAG/lightrag/utils.py", line 185, in wait_func
    result = await func(*args, **kwargs)
  File "/Users/konradwojciechowski/PyCharmProjects/graph-rag/LightRAG/lightrag/utils.py", line 108, in __call__
    return await self.func(*args, **kwargs)
  File "/Users/konradwojciechowski/PyCharmProjects/graph-rag/LightRAG/lightrag/llm/hf.py", line 146, in hf_embed
    outputs = embed_model(input_ids)
  File "/Users/konradwojciechowski/PyCharmProjects/graph-rag/LightRAG/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/Users/konradwojciechowski/PyCharmProjects/graph-rag/LightRAG/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/konradwojciechowski/.cache/huggingface/modules/transformers_modules/nomic-ai/nomic-bert-2048/359596ab182dcf943b7ca9e3f8809b6c2eaf652f/modeling_hf_nomic_bert.py", line 1909, in forward
    attention_mask = self.get_extended_attention_mask(attention_mask, hidden_states.shape[:-1])
  File "/Users/konradwojciechowski/PyCharmProjects/graph-rag/LightRAG/.venv/lib/python3.10/site-packages/transformers/modeling_utils.py", line 1069, in get_extended_attention_mask
    if not (attention_mask.dim() == 2 and self.config.is_decoder):
AttributeError: 'NoneType' object has no attribute 'dim'

You can reproduce this by changing mini-lm to nomic. Here's my script, other things I changed like llm is irrelevant here.

import os

from lightrag import LightRAG, QueryParam
from lightrag.llm.hf import hf_embed
from lightrag.llm.ollama import ollama_model_complete
from lightrag.utils import EmbeddingFunc
from transformers import AutoModel, AutoTokenizer

WORKING_DIR = "./dickens"

if not os.path.exists(WORKING_DIR):
    os.mkdir(WORKING_DIR)

rag = LightRAG(
    working_dir=WORKING_DIR,
    # llm_model_func=hf_model_complete,
    # llm_model_name="meta-llama/Llama-3.1-8B-Instruct",
    llm_model_func=ollama_model_complete,
    llm_model_name="qwen2.5:1.5b",
    llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 32768}},
    llm_model_max_async=4,
    llm_model_max_token_size=32768,
    embedding_func=EmbeddingFunc(
        # embedding_dim=384,
        embedding_dim=768,
        # max_token_size=5000,
        max_token_size=8192,
        func=lambda texts: hf_embed(
            texts,
            tokenizer=AutoTokenizer.from_pretrained(
                # "sentence-transformers/all-MiniLM-L6-v2",
                "nomic-ai/nomic-embed-text-v1.5",
                trust_remote_code=True,
                device="cpu",
            ),
            embed_model=AutoModel.from_pretrained(
                # "sentence-transformers/all-MiniLM-L6-v2",
                "nomic-ai/nomic-embed-text-v1.5",
                trust_remote_code=True,
                device="cpu",
            ),
        ),
    ),
)


with open("./book.txt", "r", encoding="utf-8") as f:
    rag.insert(f.read())

# Perform naive search
print(
    rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
)

# Perform local search
print(
    rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
)

# Perform global search
print(
    rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
)

# Perform hybrid search
print(
    rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
)

Related Issues

Not that I'm aware of.

Changes Made

added attention mask in hf.hf_embed()

Checklist

• Changes tested locally
• Code reviewed
• Documentation updated (if necessary)
• Unit tests added (if applicable)

Additional Notes

[Add any additional notes or context for the reviewer(s).]

[LarFii]

Thanks for sharing!