mirror of
https://github.com/ggerganov/llama.cpp
synced 2026-04-09 00:35:42 +02:00
f17b3be63f
* llama : fix pooling assertion crash in chunked GDN detection path
The chunked fused Gated Delta Net detection in sched_reserve() calls
graph_reserve(16*n_seqs, n_seqs, n_outputs, ...) where n_outputs = n_seqs.
This creates a dimension mismatch in build_pooling() for embedding models
with mean/rank pooling: build_inp_mean() creates a tensor with shape
[n_tokens=16*n_seqs, ...] while t_embd is reduced to [n_outputs=n_seqs, ...]
via out_ids, causing ggml_mul_mat to assert on ggml_can_mul_mat(a, b).
Fix: pass n_tokens as n_outputs in the chunked GDN graph reservation,
matching the pattern used by the pp/tg worst-case reservations.
Regression introduced by #20340 (d28961d).
Same class of bug as #12517, fixed by #12545.
* server : add mean pooling tests to embedding test suite
Add test_embedding_pooling_mean and test_embedding_pooling_mean_multiple
to cover the --pooling mean codepath, which was previously untested.
These tests would have caught the regression introduced by #20340 where
build_pooling() crashes with a ggml_mul_mat assertion due to mismatched
dimensions in the chunked GDN detection path.
---------
Co-authored-by: Domenico Crupi <domenico@zerovolt.it>
292 lines
8.8 KiB
Python
292 lines
8.8 KiB
Python
import base64
|
|
import struct
|
|
import pytest
|
|
from openai import OpenAI
|
|
from utils import *
|
|
|
|
server = ServerPreset.bert_bge_small()
|
|
|
|
EPSILON = 1e-3
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def create_server():
|
|
global server
|
|
server = ServerPreset.bert_bge_small()
|
|
|
|
|
|
def test_embedding_single():
|
|
global server
|
|
server.pooling = 'last'
|
|
server.start()
|
|
res = server.make_request("POST", "/v1/embeddings", data={
|
|
"input": "I believe the meaning of life is",
|
|
})
|
|
assert res.status_code == 200
|
|
assert len(res.body['data']) == 1
|
|
assert 'embedding' in res.body['data'][0]
|
|
assert len(res.body['data'][0]['embedding']) > 1
|
|
|
|
# make sure embedding vector is normalized
|
|
assert abs(sum([x ** 2 for x in res.body['data'][0]['embedding']]) - 1) < EPSILON
|
|
|
|
|
|
def test_embedding_multiple():
|
|
global server
|
|
server.pooling = 'last'
|
|
server.start()
|
|
res = server.make_request("POST", "/v1/embeddings", data={
|
|
"input": [
|
|
"I believe the meaning of life is",
|
|
"Write a joke about AI from a very long prompt which will not be truncated",
|
|
"This is a test",
|
|
"This is another test",
|
|
],
|
|
})
|
|
assert res.status_code == 200
|
|
assert len(res.body['data']) == 4
|
|
for d in res.body['data']:
|
|
assert 'embedding' in d
|
|
assert len(d['embedding']) > 1
|
|
|
|
|
|
def test_embedding_multiple_with_fa():
|
|
server = ServerPreset.bert_bge_small_with_fa()
|
|
server.pooling = 'last'
|
|
server.start()
|
|
# one of these should trigger the FA branch (i.e. context size % 256 == 0)
|
|
res = server.make_request("POST", "/v1/embeddings", data={
|
|
"input": [
|
|
"a "*253,
|
|
"b "*254,
|
|
"c "*255,
|
|
"d "*256,
|
|
],
|
|
})
|
|
assert res.status_code == 200
|
|
assert len(res.body['data']) == 4
|
|
for d in res.body['data']:
|
|
assert 'embedding' in d
|
|
assert len(d['embedding']) > 1
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input,is_multi_prompt",
|
|
[
|
|
# do not crash on empty input
|
|
("", False),
|
|
# single prompt
|
|
("string", False),
|
|
([12, 34, 56], False),
|
|
([12, 34, "string", 56, 78], False),
|
|
# multiple prompts
|
|
(["string1", "string2"], True),
|
|
(["string1", [12, 34, 56]], True),
|
|
([[12, 34, 56], [12, 34, 56]], True),
|
|
([[12, 34, 56], [12, "string", 34, 56]], True),
|
|
]
|
|
)
|
|
def test_embedding_mixed_input(input, is_multi_prompt: bool):
|
|
global server
|
|
server.start()
|
|
res = server.make_request("POST", "/v1/embeddings", data={"input": input})
|
|
assert res.status_code == 200
|
|
data = res.body['data']
|
|
if is_multi_prompt:
|
|
assert len(data) == len(input)
|
|
for d in data:
|
|
assert 'embedding' in d
|
|
assert len(d['embedding']) > 1
|
|
else:
|
|
assert 'embedding' in data[0]
|
|
assert len(data[0]['embedding']) > 1
|
|
|
|
|
|
def test_embedding_pooling_mean():
|
|
global server
|
|
server.pooling = 'mean'
|
|
server.start()
|
|
res = server.make_request("POST", "/v1/embeddings", data={
|
|
"input": "I believe the meaning of life is",
|
|
})
|
|
assert res.status_code == 200
|
|
assert len(res.body['data']) == 1
|
|
assert 'embedding' in res.body['data'][0]
|
|
assert len(res.body['data'][0]['embedding']) > 1
|
|
|
|
# make sure embedding vector is normalized
|
|
assert abs(sum([x ** 2 for x in res.body['data'][0]['embedding']]) - 1) < EPSILON
|
|
|
|
|
|
def test_embedding_pooling_mean_multiple():
|
|
global server
|
|
server.pooling = 'mean'
|
|
server.start()
|
|
res = server.make_request("POST", "/v1/embeddings", data={
|
|
"input": [
|
|
"I believe the meaning of life is",
|
|
"Write a joke about AI",
|
|
"This is a test",
|
|
],
|
|
})
|
|
assert res.status_code == 200
|
|
assert len(res.body['data']) == 3
|
|
for d in res.body['data']:
|
|
assert 'embedding' in d
|
|
assert len(d['embedding']) > 1
|
|
|
|
|
|
def test_embedding_pooling_none():
|
|
global server
|
|
server.pooling = 'none'
|
|
server.start()
|
|
res = server.make_request("POST", "/embeddings", data={
|
|
"input": "hello hello hello",
|
|
})
|
|
assert res.status_code == 200
|
|
assert 'embedding' in res.body[0]
|
|
assert len(res.body[0]['embedding']) == 5 # 3 text tokens + 2 special
|
|
|
|
# make sure embedding vector is not normalized
|
|
for x in res.body[0]['embedding']:
|
|
assert abs(sum([x ** 2 for x in x]) - 1) > EPSILON
|
|
|
|
|
|
def test_embedding_pooling_none_oai():
|
|
global server
|
|
server.pooling = 'none'
|
|
server.start()
|
|
res = server.make_request("POST", "/v1/embeddings", data={
|
|
"input": "hello hello hello",
|
|
})
|
|
|
|
# /v1/embeddings does not support pooling type 'none'
|
|
assert res.status_code == 400
|
|
assert "error" in res.body
|
|
|
|
|
|
def test_embedding_openai_library_single():
|
|
global server
|
|
server.pooling = 'last'
|
|
server.start()
|
|
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
|
|
res = client.embeddings.create(model="text-embedding-3-small", input="I believe the meaning of life is")
|
|
assert len(res.data) == 1
|
|
assert len(res.data[0].embedding) > 1
|
|
|
|
|
|
def test_embedding_openai_library_multiple():
|
|
global server
|
|
server.pooling = 'last'
|
|
server.start()
|
|
client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")
|
|
res = client.embeddings.create(model="text-embedding-3-small", input=[
|
|
"I believe the meaning of life is",
|
|
"Write a joke about AI from a very long prompt which will not be truncated",
|
|
"This is a test",
|
|
"This is another test",
|
|
])
|
|
assert len(res.data) == 4
|
|
for d in res.data:
|
|
assert len(d.embedding) > 1
|
|
|
|
|
|
def test_embedding_error_prompt_too_long():
|
|
global server
|
|
server.pooling = 'last'
|
|
server.start()
|
|
res = server.make_request("POST", "/v1/embeddings", data={
|
|
"input": "This is a test " * 512,
|
|
})
|
|
assert res.status_code != 200
|
|
assert "too large" in res.body["error"]["message"]
|
|
|
|
|
|
def test_same_prompt_give_same_result():
|
|
server.pooling = 'last'
|
|
server.start()
|
|
res = server.make_request("POST", "/v1/embeddings", data={
|
|
"input": [
|
|
"I believe the meaning of life is",
|
|
"I believe the meaning of life is",
|
|
"I believe the meaning of life is",
|
|
"I believe the meaning of life is",
|
|
"I believe the meaning of life is",
|
|
],
|
|
})
|
|
assert res.status_code == 200
|
|
assert len(res.body['data']) == 5
|
|
for i in range(1, len(res.body['data'])):
|
|
v0 = res.body['data'][0]['embedding']
|
|
vi = res.body['data'][i]['embedding']
|
|
for x, y in zip(v0, vi):
|
|
assert abs(x - y) < EPSILON
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"content,n_tokens",
|
|
[
|
|
("I believe the meaning of life is", 9),
|
|
("This is a test", 6),
|
|
]
|
|
)
|
|
def test_embedding_usage_single(content, n_tokens):
|
|
global server
|
|
server.start()
|
|
res = server.make_request("POST", "/v1/embeddings", data={"input": content})
|
|
assert res.status_code == 200
|
|
assert res.body['usage']['prompt_tokens'] == res.body['usage']['total_tokens']
|
|
assert res.body['usage']['prompt_tokens'] == n_tokens
|
|
|
|
|
|
def test_embedding_usage_multiple():
|
|
global server
|
|
server.start()
|
|
res = server.make_request("POST", "/v1/embeddings", data={
|
|
"input": [
|
|
"I believe the meaning of life is",
|
|
"I believe the meaning of life is",
|
|
],
|
|
})
|
|
assert res.status_code == 200
|
|
assert res.body['usage']['prompt_tokens'] == res.body['usage']['total_tokens']
|
|
assert res.body['usage']['prompt_tokens'] == 2 * 9
|
|
|
|
|
|
def test_embedding_openai_library_base64():
|
|
server.start()
|
|
test_input = "Test base64 embedding output"
|
|
|
|
# get embedding in default format
|
|
res = server.make_request("POST", "/v1/embeddings", data={
|
|
"input": test_input
|
|
})
|
|
assert res.status_code == 200
|
|
vec0 = res.body["data"][0]["embedding"]
|
|
|
|
# get embedding in base64 format
|
|
res = server.make_request("POST", "/v1/embeddings", data={
|
|
"input": test_input,
|
|
"encoding_format": "base64"
|
|
})
|
|
|
|
assert res.status_code == 200
|
|
assert "data" in res.body
|
|
assert len(res.body["data"]) == 1
|
|
|
|
embedding_data = res.body["data"][0]
|
|
assert "embedding" in embedding_data
|
|
assert isinstance(embedding_data["embedding"], str)
|
|
|
|
# Verify embedding is valid base64
|
|
decoded = base64.b64decode(embedding_data["embedding"])
|
|
# Verify decoded data can be converted back to float array
|
|
float_count = len(decoded) // 4 # 4 bytes per float
|
|
floats = struct.unpack(f'{float_count}f', decoded)
|
|
assert len(floats) > 0
|
|
assert all(isinstance(x, float) for x in floats)
|
|
assert len(floats) == len(vec0)
|
|
|
|
# make sure the decoded data is the same as the original
|
|
for x, y in zip(floats, vec0):
|
|
assert abs(x - y) < EPSILON
|