abetlen / abetlen/llama-cpp-python

`Llama.embed` crashes when `n_batch` > 512

オープン
#1,762 コメント 3 件 リアクション 0 件 担当者 0 名 GitHub で見る
bug
主要言語
Python
スター
10.6k
フォーク
1.4k
PR マージ指標
PR 指標を取得中

説明

# Expected Behavior

Embedding text with a long-context model like BGE-M3 [1] should be able to output token embeddings for more than 512 tokens (this is of interest for 'late interaction' retrieval [2]).

Llama-cpp-python will truncate the input tokens to the first n_batch tokens, where n_batch is 512 by default. The expected behaviour is that setting n_batch to a larger value would allow computing the token embeddings for longer sequences.

[1] https://huggingface.co/BAAI/bge-m3
[2] https://jina.ai/news/what-is-colbert-and-late-interaction-and-why-they-matter-in-search/

# Current Behavior

The kernel crashes when embedding text with `n_batch` > 512. This crash is not specific to the embedding model, for a few models I've tried.

# Steps to Reproduce

On a Google Colab T4 instance:

```python
%pip install --quiet --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/12.2 llama-cpp-python==0.3.0

from llama_cpp import LLAMA_POOLING_TYPE_NONE, Llama

embedder = Llama.from_pretrained(
repo_id="lm-kit/bge-m3-gguf",
filename="*F16.gguf",
n_ctx=0, # Model context is 8192
n_gpu_layers=-1,
n_batch=513, # ← Any value larger than 512 (the default) causes a crash
embedding=True,
pooling_type=LLAMA_POOLING_TYPE_NONE,
verbose=False
)

text = "Hello world" * 1000
embedding = embedder.embed(text) # ← Crash 💥
len(embedding)
```

コントリビューションガイド

コントリビューションガイドを開く

評価

この issue はまだ評価されていません。

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。