Skip to content

Commit

Permalink
server : fill usage info in reranking response
Browse files Browse the repository at this point in the history
  • Loading branch information
krystiancha committed Dec 17, 2024
1 parent 9e87358 commit 0924497
Show file tree
Hide file tree
Showing 3 changed files with 35 additions and 5 deletions.
8 changes: 6 additions & 2 deletions examples/server/server.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -738,14 +738,17 @@ struct server_task_result_rerank : server_task_result {
int index = 0;
float score = -1e6;

int32_t n_prompt_tokens;

virtual int get_index() override {
return index;
}

virtual json to_json() override {
return json {
{"index", index},
{"score", score},
{"index", index},
{"score", score},
{"tokens_evaluated", n_prompt_tokens},
};
}
};
Expand Down Expand Up @@ -2034,6 +2037,7 @@ struct server_context {
auto res = std::make_unique<server_task_result_rerank>();
res->id = slot.id_task;
res->index = slot.index;
res->n_prompt_tokens = slot.n_prompt_tokens;

for (int i = 0; i < batch.n_tokens; ++i) {
if (!batch.logits[i] || batch.seq_id[i][0] != slot.id) {
Expand Down
23 changes: 23 additions & 0 deletions examples/server/tests/unit/test_rerank.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,3 +53,26 @@ def test_invalid_rerank_req(documents):
})
assert res.status_code == 400
assert "error" in res.body


@pytest.mark.parametrize(
"query,doc1,doc2,n_tokens",
[
("Machine learning is", "A machine", "Learning is", 19),
("Which city?", "Machine learning is ", "Paris, capitale de la", 26),
]
)
def test_rerank_usage(query, doc1, doc2, n_tokens):
global server
server.start()

res = server.make_request("POST", "/rerank", data={
"query": query,
"documents": [
doc1,
doc2,
]
})
assert res.status_code == 200
assert res.body['usage']['prompt_tokens'] == res.body['usage']['total_tokens']
assert res.body['usage']['prompt_tokens'] == n_tokens
9 changes: 6 additions & 3 deletions examples/server/utils.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -587,20 +587,23 @@ static json format_embeddings_response_oaicompat(const json & request, const jso

static json format_response_rerank(const json & request, const json & ranks) {
json data = json::array();
int32_t n_prompt_tokens = 0;
int i = 0;
for (const auto & rank : ranks) {
data.push_back(json{
{"index", i++},
{"relevance_score", json_value(rank, "score", 0.0)},
});

n_prompt_tokens += json_value(rank, "tokens_evaluated", 0);
}

json res = json {
{"model", json_value(request, "model", std::string(DEFAULT_OAICOMPAT_MODEL))},
{"object", "list"},
{"usage", json { // TODO: fill
{"prompt_tokens", 0},
{"total_tokens", 0}
{"usage", json {
{"prompt_tokens", n_prompt_tokens},
{"total_tokens", n_prompt_tokens}
}},
{"results", data}
};
Expand Down

0 comments on commit 0924497

Please sign in to comment.