add_to_index uses too much GPU RAM and crashes about ragatouille

[dineshdharme]

I am facing this issue even with just 100 documents trying to use just RAG.index().
Here's a simple code which reproduces the issue on 16GB VRAM.

from ragatouille import RAGPretrainedModel
from datasets import load_dataset

another_dataset = "kdawoud91/News_Articles"
df = load_dataset(another_dataset, cache_dir=cache_hf_dir)['train'].to_pandas()

df_limit = df[:100]
df_limit = df_limit.reset_index()
df_limit = df_limit.rename(columns={"index": "id"})
df_limit["id"] = df_limit["id"].map(repr)

list_news_articles = df_limit.text.values

RAG = RAGPretrainedModel.from_pretrained("jinaai/jina-colbert-v1-en")


RAG.index(
    collection=list_news_articles,
    index_name=INDEX_NAME,
    split_documents=True,
    use_faiss=True,
    max_document_length=4060,
)
 

[bclavie]

Hey, apologies for the very late answer. This seems to be an issue with the torch k-means implementation being far too hungry in some cases. Working on a fix now.