abetlen / abetlen/llama-cpp-python

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

Đang mở
#1,762 3 bình luận 0 reaction 0 người được giao Xem trên GitHub
bug
Ngôn ngữ chính
Python
Star
10.6k
Fork
1.4k
Chỉ số merge pull request
Chỉ số pull request đang chờ

Mô tả

# 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)
```

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.