Skip to content

GitLab

Unverified Commit 62f91f78 authored by OlivierDehaene's avatar OlivierDehaene Committed by GitHub
Browse files

feat(server): support vectorized warpers in flash causal lm (#317) 


Co-authored-by: default avatarJoel Lamy-Poirier <joel.lamy-poirier@servicenow.com>
parent 951930fb
Hide whitespace changes
Inline Side-by-side
Showing with 519 additions and 232 deletions
integration-tests/models/__snapshots__/test_bloom_560m/test_bloom_560m_all_params.json
View file @ 62f91f78
...@@ -34,65 +34,65 @@ ...@@ -34,65 +34,65 @@
"tokens": [ "tokens": [
{ {
"id": 408, "id": 408,
"logprob": -1.9267578, "logprob": -0.07891846,
"special": false, "special": false,
"text": " que" "text": " que"
}, },
{ {
"id": 20288, "id": 366,
"logprob": -2.9257812, "logprob": -1.2939453,
"special": false, "special": false,
"text": " l'on" "text": " la"
}, },
{ {
"id": 22255, "id": 8769,
"logprob": -2.8964844, "logprob": -0.3708496,
"special": false, "special": false,
"text": " trouve" "text": " personne"
}, },
{ {
"id": 1622, "id": 1479,
"logprob": -1.1083984, "logprob": -2.2871094,
"special": false, "special": false,
"text": " une" "text": " qui"
}, },
{ {
"id": 187079, "id": 2997,
"logprob": -7.796875, "logprob": -0.8671875,
"special": false, "special": false,
"text": " posture" "text": " vous"
}, },
{ {
"id": 501, "id": 35977,
"logprob": -5.390625, "logprob": -1.5097656,
"special": false, "special": false,
"text": " par" "text": " suit"
}, },
{ {
"id": 8741, "id": 21558,
"logprob": -0.34936523, "logprob": -0.07891846,
"special": false, "special": false,
"text": " rapport" "text": " ait"
}, },
{ {
"id": 693, "id": 447,
"logprob": 0.0, "logprob": -0.12695312,
"special": false, "special": false,
"text": " à" "text": " un"
}, },
{ {
"id": 366, "id": 78606,
"logprob": -2.3378906, "logprob": -2.21875,
"special": false, "special": false,
"text": " la" "text": " profil"
}, },
{ {
"id": 36503, "id": 3899,
"logprob": -3.6640625, "logprob": -1.3535156,
"special": false, "special": false,
"text": " pratique" "text": " bien"
} }
] ]
}, },
"generated_text": "Pour déguster un ortolan, il faut tout d'abord que l'on trouve une posture par rapport à la pratique" "generated_text": "Pour déguster un ortolan, il faut tout d'abord que la personne qui vous suit ait un profil bien"
} }
integration-tests/models/__snapshots__/test_flash_llama/test_flash_llama_all_params.json
View file @ 62f91f78
{ {
"details": { "details": {
"best_of_sequences": null, "best_of_sequences": null,
"finish_reason": "length", "finish_reason": "stop_sequence",
"generated_tokens": 10, "generated_tokens": 5,
"prefill": [ "prefill": [
{ {
"id": 1, "id": 1,
...@@ -24,65 +24,35 @@ ...@@ -24,65 +24,35 @@
"tokens": [ "tokens": [
{ {
"id": 5229, "id": 5229,
"logprob": -3.3085938, "logprob": -2.5683594,
"special": false, "special": false,
"text": " failed" "text": " failed"
}, },
{
"id": 363,
"logprob": -3.984375,
"special": false,
"text": " for"
},
{
"id": 5641,
"logprob": -6.53125,
"special": false,
"text": " IP"
},
{
"id": 16428,
"logprob": -3.1835938,
"special": false,
"text": " Address"
},
{ {
"id": 29901, "id": 29901,
"logprob": -1.2324219, "logprob": -0.45336914,
"special": false, "special": false,
"text": ":" "text": ":"
}, },
{ {
"id": 525, "id": 4829,
"logprob": -2.6855469, "logprob": -1.8408203,
"special": false,
"text": " '"
},
{
"id": 8516,
"logprob": -7.1601562,
"special": false,
"text": "None"
},
{
"id": 4286,
"logprob": -2.4433594,
"special": false, "special": false,
"text": "'." "text": " Error"
}, },
{ {
"id": 13, "id": 297,
"logprob": -0.06530762, "logprob": -1.0556641,
"special": false, "special": false,
"text": "\n" "text": " in"
}, },
{ {
"id": 294, "id": 1243,
"logprob": -7.953125, "logprob": 0.0,
"special": false, "special": false,
"text": "as" "text": " test"
} }
] ]
}, },
"generated_text": "Test requestfailed for IP Address: 'None'.\nas" "generated_text": "Test requestfailed: Error in test"
} }
integration-tests/models/__snapshots__/test_flash_starcoder/test_flash_starcoder_default_params.json
View file @ 62f91f78
{ {
"details": { "details": {
"best_of_sequences": null, "best_of_sequences": null,
"finish_reason": "eos_token", "finish_reason": "length",
"generated_tokens": 12, "generated_tokens": 60,
"prefill": [ "prefill": [
{ {
"id": 589, "id": 589,
...@@ -29,77 +29,365 @@ ...@@ -29,77 +29,365 @@
"tokens": [ "tokens": [
{ {
"id": 2262, "id": 2262,
"logprob": -0.7451172, "logprob": -0.042999268,
"special": false, "special": false,
"text": "():" "text": "():"
}, },
{ {
"id": 284, "id": 284,
"logprob": -0.21325684, "logprob": 0.0,
"special": false, "special": false,
"text": "\n " "text": "\n "
}, },
{ {
"id": 5741, "id": 1459,
"logprob": -5.734375, "logprob": 0.0,
"special": false,
"text": " print"
},
{
"id": 440,
"logprob": 0.0,
"special": false,
"text": "(\""
},
{
"id": 8279,
"logprob": 0.0,
"special": false,
"text": "Hello"
},
{
"id": 10896,
"logprob": -0.3659668,
"special": false,
"text": " World"
},
{
"id": 657,
"logprob": -0.49804688,
"special": false,
"text": "\")"
},
{
"id": 203,
"logprob": -0.11279297,
"special": false, "special": false,
"text": " logging" "text": "\n"
}, },
{ {
"id": 32, "id": 203,
"logprob": 0.0, "logprob": 0.0,
"special": false, "special": false,
"text": "." "text": "\n"
}, },
{ {
"id": 1338, "id": 589,
"logprob": -0.3232422, "logprob": -0.20141602,
"special": false, "special": false,
"text": "info" "text": "def"
}, },
{ {
"id": 463, "id": 1459,
"logprob": -1.0380859, "logprob": 0.0,
"special": false,
"text": " print"
},
{
"id": 81,
"logprob": 0.0,
"special": false,
"text": "_"
},
{
"id": 7656,
"logprob": 0.0,
"special": false,
"text": "hello"
},
{
"id": 81,
"logprob": 0.0,
"special": false,
"text": "_"
},
{
"id": 426,
"logprob": -0.051635742,
"special": false,
"text": "name"
},
{
"id": 26,
"logprob": 0.0,
"special": false,
"text": "("
},
{
"id": 426,
"logprob": 0.0,
"special": false, "special": false,
"text": "('" "text": "name"
},
{
"id": 711,
"logprob": 0.0,
"special": false,
"text": "):"
},
{
"id": 284,
"logprob": 0.0,
"special": false,
"text": "\n "
},
{
"id": 1459,
"logprob": 0.0,
"special": false,
"text": " print"
},
{
"id": 440,
"logprob": -0.16027832,
"special": false,
"text": "(\""
}, },
{ {
"id": 8279, "id": 8279,
"logprob": -0.8378906, "logprob": 0.0,
"special": false, "special": false,
"text": "Hello" "text": "Hello"
}, },
{
"id": 313,
"logprob": 0.0,
"special": false,
"text": " \""
},
{
"id": 474,
"logprob": 0.0,
"special": false,
"text": " +"
},
{
"id": 636,
"logprob": 0.0,
"special": false,
"text": " name"
},
{
"id": 27,
"logprob": 0.0,
"special": false,
"text": ")"
},
{
"id": 203,
"logprob": 0.0,
"special": false,
"text": "\n"
},
{
"id": 203,
"logprob": 0.0,
"special": false,
"text": "\n"
},
{
"id": 589,
"logprob": 0.0,
"special": false,
"text": "def"
},
{
"id": 1459,
"logprob": 0.0,
"special": false,
"text": " print"
},
{
"id": 81,
"logprob": 0.0,
"special": false,
"text": "_"
},
{
"id": 7656,
"logprob": 0.0,
"special": false,
"text": "hello"
},
{
"id": 81,
"logprob": 0.0,
"special": false,
"text": "_"
},
{
"id": 426,
"logprob": 0.0,
"special": false,
"text": "name"
},
{
"id": 81,
"logprob": 0.0,
"special": false,
"text": "_"
},
{
"id": 381,
"logprob": 0.0,
"special": false,
"text": "age"
},
{
"id": 26,
"logprob": 0.0,
"special": false,
"text": "("
},
{
"id": 426,
"logprob": 0.0,
"special": false,
"text": "name"
},
{ {
"id": 30, "id": 30,
"logprob": -1.9501953, "logprob": 0.0,
"special": false, "special": false,
"text": "," "text": ","
}, },
{ {
"id": 10896, "id": 11442,
"logprob": -1.3476562, "logprob": 0.0,
"special": false, "special": false,
"text": " World" "text": " age"
},
{
"id": 711,
"logprob": 0.0,
"special": false,
"text": "):"
},
{
"id": 284,
"logprob": 0.0,
"special": false,
"text": "\n "
},
{
"id": 1459,
"logprob": 0.0,
"special": false,
"text": " print"
},
{
"id": 440,
"logprob": 0.0,
"special": false,
"text": "(\""
}, },
{ {
"id": 683, "id": 8279,
"logprob": -1.796875, "logprob": 0.0,
"special": false,
"text": "Hello"
},
{
"id": 313,
"logprob": 0.0,
"special": false,
"text": " \""
},
{
"id": 474,
"logprob": 0.0,
"special": false,
"text": " +"
},
{
"id": 636,
"logprob": 0.0,
"special": false,
"text": " name"
},
{
"id": 474,
"logprob": 0.0,
"special": false,
"text": " +"
},
{
"id": 313,
"logprob": -0.6328125,
"special": false,
"text": " \""
},
{
"id": 313,
"logprob": -1.7011719,
"special": false,
"text": " \""
},
{
"id": 474,
"logprob": 0.0,
"special": false,
"text": " +"
},
{
"id": 596,
"logprob": 0.0,
"special": false,
"text": " str"
},
{
"id": 26,
"logprob": 0.0,
"special": false,
"text": "("
},
{
"id": 381,
"logprob": 0.0,
"special": false,
"text": "age"
},
{
"id": 490,
"logprob": 0.0,
"special": false, "special": false,
"text": "')" "text": "))"
}, },
{ {
"id": 203, "id": 203,
"logprob": -0.9873047, "logprob": 0.0,
"special": false, "special": false,
"text": "\n" "text": "\n"
}, },
{ {
"id": 0, "id": 203,
"logprob": -0.7495117, "logprob": 0.0,
"special": true, "special": false,
"text": "<|endoftext|>" "text": "\n"
},
{
"id": 589,
"logprob": 0.0,
"special": false,
"text": "def"
},
{
"id": 1459,
"logprob": 0.0,
"special": false,
"text": " print"
} }
] ]
}, },
"generated_text": "():\n logging.info('Hello, World')\n<|endoftext|>" "generated_text": "():\n print(\"Hello World\")\n\ndef print_hello_name(name):\n print(\"Hello \" + name)\n\ndef print_hello_name_age(name, age):\n print(\"Hello \" + name + \" \" + str(age))\n\ndef print"
} }
integration-tests/models/__snapshots__/test_mt0_base/test_mt0_base_all_params.json
View file @ 62f91f78
{ {
"details": { "details": {
"best_of_sequences": null, "best_of_sequences": null,
"finish_reason": "length", "finish_reason": "eos_token",
"generated_tokens": 10, "generated_tokens": 9,
"prefill": [ "prefill": [
{ {
"id": 0, "id": 0,
...@@ -14,65 +14,59 @@ ...@@ -14,65 +14,59 @@
"tokens": [ "tokens": [
{ {
"id": 16017, "id": 16017,
"logprob": -1.3505859, "logprob": -0.30908203,
"special": false, "special": false,
"text": " blue" "text": " blue"
}, },
{ {
"id": 20495, "id": 20495,
"logprob": -0.50439453, "logprob": 0.0,
"special": false, "special": false,
"text": " sky" "text": " sky"
}, },
{ {
"id": 259, "id": 259,
"logprob": -1.2011719, "logprob": -0.28271484,
"special": false, "special": false,
"text": " " "text": " "
}, },
{ {
"id": 15484, "id": 15484,
"logprob": -2.8378906, "logprob": -1.7929688,
"special": false, "special": false,
"text": "appear" "text": "appear"
}, },
{ {
"id": 345, "id": 345,
"logprob": -0.87597656, "logprob": -0.8935547,
"special": false, "special": false,
"text": "ed" "text": "ed"
}, },
{ {
"id": 288, "id": 281,
"logprob": -1.8447266, "logprob": 0.0,
"special": false, "special": false,
"text": " to" "text": " in"
}, },
{ {
"id": 35622, "id": 287,
"logprob": -7.1445312, "logprob": 0.0,
"special": false, "special": false,
"text": " cloud" "text": " the"
}, },
{ {
"id": 263, "id": 20495,
"logprob": -1.2929688, "logprob": -0.32299805,
"special": false,
"text": "s"
},
{
"id": 14701,
"logprob": -3.0761719,
"special": false, "special": false,
"text": " above" "text": " sky"
}, },
{ {
"id": 751, "id": 1,
"logprob": -4.4375, "logprob": 0.0,
"special": false, "special": true,
"text": " all" "text": "</s>"
} }
] ]
}, },
"generated_text": "Why is the sky blue?blue sky appeared to clouds above all" "generated_text": "Why is the sky blue?blue sky appeared in the sky"
} }
integration-tests/models/test_flash_llama.py
View file @ 62f91f78
...@@ -40,7 +40,7 @@ async def test_flash_llama_all_params(flash_llama, response_snapshot): ...@@ -40,7 +40,7 @@ async def test_flash_llama_all_params(flash_llama, response_snapshot):
seed=0, seed=0,
) )
assert response.details.generated_tokens == 10 assert response.details.generated_tokens == 5
assert response == response_snapshot assert response == response_snapshot
......
integration-tests/models/test_flash_starcoder.py
View file @ 62f91f78
...@@ -29,7 +29,7 @@ async def test_flash_starcoder_default_params(flash_starcoder, response_snapshot ...@@ -29,7 +29,7 @@ async def test_flash_starcoder_default_params(flash_starcoder, response_snapshot
"def print_hello", max_new_tokens=60, temperature=0.2, top_p=0.95, seed=0 "def print_hello", max_new_tokens=60, temperature=0.2, top_p=0.95, seed=0
) )
assert response.details.generated_tokens == 12 assert response.details.generated_tokens == 60
assert response == response_snapshot assert response == response_snapshot
......
integration-tests/models/test_mt0_base.py
View file @ 62f91f78
...@@ -43,7 +43,7 @@ async def test_mt0_base_all_params(mt0_base, response_snapshot): ...@@ -43,7 +43,7 @@ async def test_mt0_base_all_params(mt0_base, response_snapshot):
seed=0, seed=0,
) )
assert response.details.generated_tokens == 10 assert response.details.generated_tokens == 9
assert response == response_snapshot assert response == response_snapshot
......
server/tests/models/test_bloom.py
View file @ 62f91f78
...@@ -38,7 +38,7 @@ def default_pb_batch(default_pb_request): ...@@ -38,7 +38,7 @@ def default_pb_batch(default_pb_request):
@pytest.fixture @pytest.fixture
def default_bloom_batch(default_pb_batch, bloom_560m_tokenizer): def default_bloom_batch(default_pb_batch, bloom_560m_tokenizer):
return BloomCausalLMBatch.from_pb( return BloomCausalLMBatch.from_pb(
default_pb_batch, bloom_560m_tokenizer, torch.device("cpu") default_pb_batch, bloom_560m_tokenizer, torch.float32, torch.device("cpu")
) )
...@@ -52,7 +52,7 @@ def default_multi_requests_bloom_batch(default_pb_request, bloom_560m_tokenizer) ...@@ -52,7 +52,7 @@ def default_multi_requests_bloom_batch(default_pb_request, bloom_560m_tokenizer)
batch_pb = generate_pb2.Batch(id=0, requests=[req_0, req_1], size=2) batch_pb = generate_pb2.Batch(id=0, requests=[req_0, req_1], size=2)
return BloomCausalLMBatch.from_pb( return BloomCausalLMBatch.from_pb(
batch_pb, bloom_560m_tokenizer, torch.device("cpu") batch_pb, bloom_560m_tokenizer, torch.float32, torch.device("cpu")
) )
...@@ -286,7 +286,9 @@ def test_batch_concatenate( ...@@ -286,7 +286,9 @@ def test_batch_concatenate(
== default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens == default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
) )
next_batch = next_batch.filter([next_batch.requests[0].id, next_batch.requests[1].id]) next_batch = next_batch.filter(
[next_batch.requests[0].id, next_batch.requests[1].id]
)
for _ in range( for _ in range(
default_bloom_batch.stopping_criterias[0].max_new_tokens default_bloom_batch.stopping_criterias[0].max_new_tokens
......
server/tests/models/test_causal_lm.py
View file @ 62f91f78
...@@ -38,7 +38,9 @@ def default_pb_batch(default_pb_request): ...@@ -38,7 +38,9 @@ def default_pb_batch(default_pb_request):
@pytest.fixture @pytest.fixture
def default_causal_lm_batch(default_pb_batch, gpt2_tokenizer): def default_causal_lm_batch(default_pb_batch, gpt2_tokenizer):
return CausalLMBatch.from_pb(default_pb_batch, gpt2_tokenizer, torch.device("cpu")) return CausalLMBatch.from_pb(
default_pb_batch, gpt2_tokenizer, torch.float32, torch.device("cpu")
)
@pytest.fixture @pytest.fixture
...@@ -50,7 +52,9 @@ def default_multi_requests_causal_lm_batch(default_pb_request, gpt2_tokenizer): ...@@ -50,7 +52,9 @@ def default_multi_requests_causal_lm_batch(default_pb_request, gpt2_tokenizer):
req_1.stopping_parameters.max_new_tokens = 5 req_1.stopping_parameters.max_new_tokens = 5
batch_pb = generate_pb2.Batch(id=1, requests=[req_0, req_1], size=2) batch_pb = generate_pb2.Batch(id=1, requests=[req_0, req_1], size=2)
return CausalLMBatch.from_pb(batch_pb, gpt2_tokenizer, torch.device("cpu")) return CausalLMBatch.from_pb(
batch_pb, gpt2_tokenizer, torch.float32, torch.device("cpu")
)
def test_batch_from_pb(default_pb_batch, default_causal_lm_batch): def test_batch_from_pb(default_pb_batch, default_causal_lm_batch):
...@@ -285,7 +289,9 @@ def test_batch_concatenate( ...@@ -285,7 +289,9 @@ def test_batch_concatenate(
== default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens == default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens
) )
next_batch = next_batch.filter([next_batch.requests[0].id, next_batch.requests[1].id]) next_batch = next_batch.filter(
[next_batch.requests[0].id, next_batch.requests[1].id]
)
for _ in range( for _ in range(
default_causal_lm_batch.stopping_criterias[0].max_new_tokens default_causal_lm_batch.stopping_criterias[0].max_new_tokens
......
server/tests/models/test_santacoder.py
View file @ 62f91f78
...@@ -45,7 +45,10 @@ def default_fim_pb_batch(default_fim_pb_request): ...@@ -45,7 +45,10 @@ def default_fim_pb_batch(default_fim_pb_request):
@pytest.mark.skip @pytest.mark.skip
def test_santacoder_generate_token_completion(default_santacoder, default_pb_batch): def test_santacoder_generate_token_completion(default_santacoder, default_pb_batch):
batch = CausalLMBatch.from_pb( batch = CausalLMBatch.from_pb(
default_pb_batch, default_santacoder.tokenizer, default_santacoder.device default_pb_batch,
default_santacoder.tokenizer,
default_santacoder.dtype,
default_santacoder.device,
) )
next_batch = batch next_batch = batch
...@@ -70,7 +73,10 @@ def test_fim_santacoder_generate_token_completion( ...@@ -70,7 +73,10 @@ def test_fim_santacoder_generate_token_completion(
default_santacoder, default_fim_pb_batch default_santacoder, default_fim_pb_batch
): ):
batch = CausalLMBatch.from_pb( batch = CausalLMBatch.from_pb(
default_fim_pb_batch, default_santacoder.tokenizer, default_santacoder.device default_fim_pb_batch,
default_santacoder.tokenizer,
default_santacoder.dtype,
default_santacoder.device,
) )
next_batch = batch next_batch = batch
......
server/tests/models/test_seq2seq_lm.py
View file @ 62f91f78
...@@ -42,7 +42,7 @@ def default_pb_batch(default_pb_request): ...@@ -42,7 +42,7 @@ def default_pb_batch(default_pb_request):
@pytest.fixture @pytest.fixture
def default_seq2seq_lm_batch(default_pb_batch, mt0_small_tokenizer): def default_seq2seq_lm_batch(default_pb_batch, mt0_small_tokenizer):
return Seq2SeqLMBatch.from_pb( return Seq2SeqLMBatch.from_pb(
default_pb_batch, mt0_small_tokenizer, torch.device("cpu") default_pb_batch, mt0_small_tokenizer, torch.float32, torch.device("cpu")
) )
...@@ -55,7 +55,9 @@ def default_multi_requests_seq2seq_lm_batch(default_pb_request, mt0_small_tokeni ...@@ -55,7 +55,9 @@ def default_multi_requests_seq2seq_lm_batch(default_pb_request, mt0_small_tokeni
req_1.stopping_parameters.max_new_tokens = 5 req_1.stopping_parameters.max_new_tokens = 5
batch_pb = generate_pb2.Batch(id=0, requests=[req_0, req_1], size=2) batch_pb = generate_pb2.Batch(id=0, requests=[req_0, req_1], size=2)
return Seq2SeqLMBatch.from_pb(batch_pb, mt0_small_tokenizer, torch.device("cpu")) return Seq2SeqLMBatch.from_pb(
batch_pb, mt0_small_tokenizer, torch.float32, torch.device("cpu")
)
def test_batch_from_pb(default_pb_batch, default_seq2seq_lm_batch): def test_batch_from_pb(default_pb_batch, default_seq2seq_lm_batch):
...@@ -323,7 +325,9 @@ def test_batch_concatenate( ...@@ -323,7 +325,9 @@ def test_batch_concatenate(
) )
assert generations[2].generated_text.generated_tokens == 5 assert generations[2].generated_text.generated_tokens == 5
next_batch = next_batch.filter([next_batch.requests[0].id, next_batch.requests[1].id]) next_batch = next_batch.filter(
[next_batch.requests[0].id, next_batch.requests[1].id]
)
generations, next_batch = default_seq2seq_lm.generate_token(next_batch) generations, next_batch = default_seq2seq_lm.generate_token(next_batch)
assert next_batch is not None assert next_batch is not None
......
server/text_generation_server/models/bloom.py
View file @ 62f91f78
...@@ -39,10 +39,11 @@ class BloomCausalLMBatch(CausalLMBatch): ...@@ -39,10 +39,11 @@ class BloomCausalLMBatch(CausalLMBatch):
cls, cls,
pb: generate_pb2.Batch, pb: generate_pb2.Batch,
tokenizer: PreTrainedTokenizerBase, tokenizer: PreTrainedTokenizerBase,
dtype: torch.dtype,
device: torch.device, device: torch.device,
) -> "CausalLMBatch": ) -> "CausalLMBatch":
batch = super(BloomCausalLMBatch, cls).from_pb( batch = super(BloomCausalLMBatch, cls).from_pb(
pb=pb, tokenizer=tokenizer, device=device pb=pb, tokenizer=tokenizer, dtype=dtype, device=device
) )
batch.keys_head_dim_last = False batch.keys_head_dim_last = False
return batch return batch
......
server/text_generation_server/models/causal_lm.py
View file @ 62f91f78
...@@ -66,6 +66,7 @@ class CausalLMBatch(Batch): ...@@ -66,6 +66,7 @@ class CausalLMBatch(Batch):
cls, cls,
pb: generate_pb2.Batch, pb: generate_pb2.Batch,
tokenizer: PreTrainedTokenizerBase, tokenizer: PreTrainedTokenizerBase,
dtype: torch.dtype,
device: torch.device, device: torch.device,
) -> "CausalLMBatch": ) -> "CausalLMBatch":
inputs = [] inputs = []
......
server/text_generation_server/models/flash_causal_lm.py
View file @ 62f91f78
...@@ -18,11 +18,7 @@ from text_generation_server.models.types import ( ...@@ -18,11 +18,7 @@ from text_generation_server.models.types import (
GeneratedText, GeneratedText,
) )
from text_generation_server.pb import generate_pb2 from text_generation_server.pb import generate_pb2
from text_generation_server.utils import ( from text_generation_server.utils import StoppingCriteria, HeterogeneousNextTokenChooser
NextTokenChooser,
StoppingCriteria,
Sampling,
)
tracer = trace.get_tracer(__name__) tracer = trace.get_tracer(__name__)
...@@ -48,7 +44,7 @@ class FlashCausalLMBatch(Batch): ...@@ -48,7 +44,7 @@ class FlashCausalLMBatch(Batch):
# All tokens # All tokens
all_input_ids: List[List[int]] all_input_ids: List[List[int]]
all_input_ids_tensor: List[torch.Tensor] all_input_ids_tensor: torch.Tensor
# Lengths of all generations present in the batch # Lengths of all generations present in the batch
input_lengths: List[int] input_lengths: List[int]
...@@ -56,7 +52,7 @@ class FlashCausalLMBatch(Batch): ...@@ -56,7 +52,7 @@ class FlashCausalLMBatch(Batch):
read_offsets: List[Optional[int]] read_offsets: List[Optional[int]]
# Generation helpers # Generation helpers
next_token_choosers: List[NextTokenChooser] next_token_chooser: HeterogeneousNextTokenChooser
stopping_criterias: List[StoppingCriteria] stopping_criterias: List[StoppingCriteria]
# Maximum number of tokens this batch will grow to # Maximum number of tokens this batch will grow to
...@@ -75,6 +71,7 @@ class FlashCausalLMBatch(Batch): ...@@ -75,6 +71,7 @@ class FlashCausalLMBatch(Batch):
cls, cls,
pb: generate_pb2.Batch, pb: generate_pb2.Batch,
tokenizer: PreTrainedTokenizerBase, tokenizer: PreTrainedTokenizerBase,
dtype: torch.dtype,
device: torch.device, device: torch.device,
) -> "FlashCausalLMBatch": ) -> "FlashCausalLMBatch":
position_ids = [] position_ids = []
...@@ -87,13 +84,14 @@ class FlashCausalLMBatch(Batch): ...@@ -87,13 +84,14 @@ class FlashCausalLMBatch(Batch):
all_input_ids = [] all_input_ids = []
requests_idx_mapping = {} requests_idx_mapping = {}
next_token_choosers = [] next_token_chooser_parameters = []
stopping_criterias = [] stopping_criterias = []
# Cumulative length # Cumulative length
cumulative_length = 0 cumulative_length = 0
max_tokens = 0 max_tokens = 0
max_length = 0
# Parse batch # Parse batch
for i, r in enumerate(pb.requests): for i, r in enumerate(pb.requests):
...@@ -119,7 +117,7 @@ class FlashCausalLMBatch(Batch): ...@@ -119,7 +117,7 @@ class FlashCausalLMBatch(Batch):
# Add cumulative lengths of all previous inputs # Add cumulative lengths of all previous inputs
cu_seqlens.append(cumulative_length + input_length) cu_seqlens.append(cumulative_length + input_length)
next_token_choosers.append(NextTokenChooser.from_pb(r.parameters, device)) next_token_chooser_parameters.append(r.parameters)
stopping_criteria = StoppingCriteria.from_pb( stopping_criteria = StoppingCriteria.from_pb(
r.stopping_parameters, tokenizer r.stopping_parameters, tokenizer
...@@ -130,11 +128,26 @@ class FlashCausalLMBatch(Batch): ...@@ -130,11 +128,26 @@ class FlashCausalLMBatch(Batch):
# Update # Update
cumulative_length += input_length cumulative_length += input_length
max_tokens += input_length + max_new_tokens max_tokens += input_length + max_new_tokens
max_length = max(max_length, input_length + max_new_tokens)
next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
next_token_chooser_parameters, dtype, device
)
# Padded all_input_ids_tensor
all_input_ids_tensor = np.zeros(
(len(all_input_ids), max_length), dtype=np.int64
)
for i, input_ids in enumerate(all_input_ids):
all_input_ids_tensor[i, : len(input_ids)] = input_ids
# Create tensors on device # Create tensors on device
input_ids = torch.tensor( input_ids = torch.tensor(
np.concatenate(all_input_ids), dtype=torch.int64, device=device np.concatenate(all_input_ids), dtype=torch.int64, device=device
) )
all_input_ids_tensor = torch.tensor(
all_input_ids_tensor, dtype=torch.int64, device=device
)
position_ids = torch.tensor( position_ids = torch.tensor(
np.concatenate(position_ids), dtype=torch.int32, device=device np.concatenate(position_ids), dtype=torch.int32, device=device
) )
...@@ -154,8 +167,8 @@ class FlashCausalLMBatch(Batch): ...@@ -154,8 +167,8 @@ class FlashCausalLMBatch(Batch):
prefix_offsets=prefix_offsets, prefix_offsets=prefix_offsets,
read_offsets=read_offsets, read_offsets=read_offsets,
all_input_ids=all_input_ids, all_input_ids=all_input_ids,
all_input_ids_tensor=[], all_input_ids_tensor=all_input_ids_tensor,
next_token_choosers=next_token_choosers, next_token_chooser=next_token_chooser,
stopping_criterias=stopping_criterias, stopping_criterias=stopping_criterias,
max_tokens=max_tokens, max_tokens=max_tokens,
) )
...@@ -176,31 +189,29 @@ class FlashCausalLMBatch(Batch): ...@@ -176,31 +189,29 @@ class FlashCausalLMBatch(Batch):
# New values after filtering # New values after filtering
requests_idx_mapping = {} requests_idx_mapping = {}
input_ids = self.input_ids.new_empty(len(request_ids)) # Used to index into tensors
position_ids = self.position_ids.new_empty(len(request_ids)) indices = []
# Create on CPU to only move to GPU once instead of at every copy # Create on CPU to only move to GPU once instead of at every copy
cu_seqlens = torch.zeros(len(request_ids) + 1, dtype=torch.int32) cu_seqlens = torch.zeros(len(request_ids) + 1, dtype=torch.int32)
cu_seqlens_q = torch.arange( cu_seqlens_q = self.cu_seqlens_q[: len(request_ids) + 1]
0, len(request_ids) + 1, device=self.cu_seqlens_q.device, dtype=torch.int32
)
max_seqlen = 0 max_seqlen = 0
past_key_values = [] past_key_values = []
requests = [] requests = []
all_input_ids = [] all_input_ids = []
all_input_ids_tensor = []
input_lengths = [] input_lengths = []
prefix_offsets = [] prefix_offsets = []
read_offsets = [] read_offsets = []
next_token_choosers = []
stopping_criterias = [] stopping_criterias = []
max_tokens = 0 max_tokens = 0
for i, request_id in enumerate(request_ids): for i, request_id in enumerate(request_ids):
idx = self.requests_idx_mapping[request_id] idx = self.requests_idx_mapping[request_id]
indices.append(idx)
requests_idx_mapping[request_id] = i requests_idx_mapping[request_id] = i
requests.append(self.requests[idx]) requests.append(self.requests[idx])
...@@ -208,10 +219,6 @@ class FlashCausalLMBatch(Batch): ...@@ -208,10 +219,6 @@ class FlashCausalLMBatch(Batch):
# Get length # Get length
request_input_length = self.input_lengths[idx] request_input_length = self.input_lengths[idx]
# Copy tensors (GPU)
input_ids[i] = self.input_ids[idx]
position_ids[i] = self.position_ids[idx]
# Copy to tensor (CPU) # Copy to tensor (CPU)
cu_seqlens[i + 1] = cumulative_length + request_input_length cu_seqlens[i + 1] = cumulative_length + request_input_length
max_seqlen = max(max_seqlen, request_input_length) max_seqlen = max(max_seqlen, request_input_length)
...@@ -222,14 +229,11 @@ class FlashCausalLMBatch(Batch): ...@@ -222,14 +229,11 @@ class FlashCausalLMBatch(Batch):
) )
all_input_ids.append(self.all_input_ids[idx]) all_input_ids.append(self.all_input_ids[idx])
all_input_ids_tensor.append(self.all_input_ids_tensor[idx])
input_lengths.append(request_input_length) input_lengths.append(request_input_length)
prefix_offsets.append(self.prefix_offsets[idx]) prefix_offsets.append(self.prefix_offsets[idx])
read_offsets.append(self.read_offsets[idx]) read_offsets.append(self.read_offsets[idx])
next_token_choosers.append(self.next_token_choosers[idx])
stopping_criteria = self.stopping_criterias[idx] stopping_criteria = self.stopping_criterias[idx]
stopping_criterias.append(stopping_criteria) stopping_criterias.append(stopping_criteria)
...@@ -258,6 +262,12 @@ class FlashCausalLMBatch(Batch): ...@@ -258,6 +262,12 @@ class FlashCausalLMBatch(Batch):
# Cat all past # Cat all past
past_key_values = torch.cat(past_key_values, dim=1) past_key_values = torch.cat(past_key_values, dim=1)
# Index into tensors
input_ids = self.input_ids[indices]
position_ids = self.position_ids[indices]
all_input_ids_tensor = self.all_input_ids_tensor[indices]
next_token_chooser = self.next_token_chooser.filter(indices)
# Move to GPU now that we have the whole tensor # Move to GPU now that we have the whole tensor
cu_seqlens = cu_seqlens.to(self.cu_seqlens.device) cu_seqlens = cu_seqlens.to(self.cu_seqlens.device)
...@@ -276,7 +286,7 @@ class FlashCausalLMBatch(Batch): ...@@ -276,7 +286,7 @@ class FlashCausalLMBatch(Batch):
read_offsets=read_offsets, read_offsets=read_offsets,
all_input_ids=all_input_ids, all_input_ids=all_input_ids,
all_input_ids_tensor=all_input_ids_tensor, all_input_ids_tensor=all_input_ids_tensor,
next_token_choosers=next_token_choosers, next_token_chooser=next_token_chooser,
stopping_criterias=stopping_criterias, stopping_criterias=stopping_criterias,
max_tokens=max_tokens, max_tokens=max_tokens,
) )
...@@ -290,6 +300,7 @@ class FlashCausalLMBatch(Batch): ...@@ -290,6 +300,7 @@ class FlashCausalLMBatch(Batch):
total_batch_size = sum([len(b) for b in batches]) total_batch_size = sum([len(b) for b in batches])
dtype = batches[0].past_key_values.dtype
device = batches[0].input_ids.device device = batches[0].input_ids.device
input_ids = batches[0].input_ids.new_empty(total_batch_size) input_ids = batches[0].input_ids.new_empty(total_batch_size)
...@@ -302,19 +313,19 @@ class FlashCausalLMBatch(Batch): ...@@ -302,19 +313,19 @@ class FlashCausalLMBatch(Batch):
past_key_values = [] past_key_values = []
all_input_ids = [] all_input_ids = []
all_input_ids_tensor = []
input_lengths = [] input_lengths = []
prefix_offsets = [] prefix_offsets = []
read_offsets = [] read_offsets = []
next_token_choosers = [] next_token_chooser_parameters = []
stopping_criterias = [] stopping_criterias = []
# Cumulative length # Cumulative length
cumulative_batch_size = 0 cumulative_batch_size = 0
cumulative_length = 0 cumulative_length = 0
max_tokens = 0 max_tokens = 0
max_length = 0
for i, batch in enumerate(batches): for i, batch in enumerate(batches):
requests.extend(batch.requests) requests.extend(batch.requests)
...@@ -347,25 +358,54 @@ class FlashCausalLMBatch(Batch): ...@@ -347,25 +358,54 @@ class FlashCausalLMBatch(Batch):
) )
all_input_ids.extend(batch.all_input_ids) all_input_ids.extend(batch.all_input_ids)
all_input_ids_tensor.extend(batch.all_input_ids_tensor)
input_lengths.extend(batch.input_lengths) input_lengths.extend(batch.input_lengths)
prefix_offsets.extend(batch.prefix_offsets) prefix_offsets.extend(batch.prefix_offsets)
read_offsets.extend(batch.read_offsets) read_offsets.extend(batch.read_offsets)
next_token_choosers.extend(batch.next_token_choosers) next_token_chooser_parameters.extend([r.parameters for r in batch.requests])
stopping_criterias.extend(batch.stopping_criterias) stopping_criterias.extend(batch.stopping_criterias)
# Update # Update
cumulative_length += batch.cu_seqlens[-1] cumulative_length += batch.cu_seqlens[-1]
cumulative_batch_size += len(batch) cumulative_batch_size += len(batch)
max_tokens += batch.max_tokens max_tokens += batch.max_tokens
max_length = max(
max_length,
max(
input_length
+ stopping_criteria.max_new_tokens
- stopping_criteria.current_tokens
for input_length, stopping_criteria in zip(
batch.input_lengths, batch.stopping_criterias
)
),
)
all_input_ids_tensor = torch.zeros(
(total_batch_size, max_length), dtype=torch.int64, device=device
)
cumulative_batch_size = 0
for i, batch in enumerate(batches):
start_index = cumulative_batch_size
end_index = cumulative_batch_size + len(batch)
all_input_ids_tensor[
start_index:end_index, : batch.all_input_ids_tensor.shape[1]
] = batch.all_input_ids_tensor[:, :max_length]
cumulative_batch_size += len(batch)
# Cat past # Cat past
past_key_values = torch.cat(past_key_values, dim=1) past_key_values = torch.cat(past_key_values, dim=1)
# Create final tensor on GPU # Create final tensor on GPU
cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
next_token_chooser_parameters, dtype=dtype, device=device
)
return FlashCausalLMBatch( return FlashCausalLMBatch(
batch_id=batches[0].batch_id, batch_id=batches[0].batch_id,
requests=requests, requests=requests,
...@@ -381,7 +421,7 @@ class FlashCausalLMBatch(Batch): ...@@ -381,7 +421,7 @@ class FlashCausalLMBatch(Batch):
read_offsets=read_offsets, read_offsets=read_offsets,
all_input_ids=all_input_ids, all_input_ids=all_input_ids,
all_input_ids_tensor=all_input_ids_tensor, all_input_ids_tensor=all_input_ids_tensor,
next_token_choosers=next_token_choosers, next_token_chooser=next_token_chooser,
stopping_criterias=stopping_criterias, stopping_criterias=stopping_criterias,
max_tokens=max_tokens, max_tokens=max_tokens,
) )
...@@ -463,6 +503,7 @@ class FlashCausalLM(Model): ...@@ -463,6 +503,7 @@ class FlashCausalLM(Model):
self, batch: FlashCausalLMBatch self, batch: FlashCausalLMBatch
) -> Tuple[List[Generation], Optional[FlashCausalLMBatch]]: ) -> Tuple[List[Generation], Optional[FlashCausalLMBatch]]:
prefill = batch.past_key_values is None prefill = batch.past_key_values is None
single_request = len(batch) == 1
if prefill and len(batch) == 1: if prefill and len(batch) == 1:
# Ask to pre-allocate kv to its max size # Ask to pre-allocate kv to its max size
...@@ -483,6 +524,17 @@ class FlashCausalLM(Model): ...@@ -483,6 +524,17 @@ class FlashCausalLM(Model):
pre_allocate_past_size, pre_allocate_past_size,
) )
if prefill:
next_token_logits = (
out[-1:] if single_request else out[batch.cu_seqlens[1:] - 1]
)
else:
next_token_logits = out
next_input_ids, next_token_logprobs = batch.next_token_chooser(
batch.all_input_ids_tensor[:, : batch.max_seqlen], next_token_logits
)
if prefill: if prefill:
if len(batch) > 1: if len(batch) > 1:
# We create the prefill_tokens_indices tensor that will be used to gather prefill logprobs # We create the prefill_tokens_indices tensor that will be used to gather prefill logprobs
...@@ -493,15 +545,11 @@ class FlashCausalLM(Model): ...@@ -493,15 +545,11 @@ class FlashCausalLM(Model):
batch.cu_seqlens_q = torch.arange( batch.cu_seqlens_q = torch.arange(
0, len(batch) + 1, device=self.device, dtype=torch.int32 0, len(batch) + 1, device=self.device, dtype=torch.int32
) )
next_input_ids = batch.input_ids.new_empty(len(batch))
next_position_ids = batch.position_ids.new_empty(len(batch)) next_position_ids = batch.position_ids.new_empty(len(batch))
else: else:
prefill_logprobs = None prefill_logprobs = None
next_input_ids = batch.input_ids
next_position_ids = batch.position_ids next_position_ids = batch.position_ids
next_token_logprobs = out.new_empty(len(batch))
# Prepare past for next decode # Prepare past for next decode
if len(batch) > 1: if len(batch) > 1:
# Used to slice next batch past # Used to slice next batch past
...@@ -552,7 +600,6 @@ class FlashCausalLM(Model): ...@@ -552,7 +600,6 @@ class FlashCausalLM(Model):
# Zipped iterator # Zipped iterator
iterator = zip( iterator = zip(
batch.input_lengths, batch.input_lengths,
batch.next_token_choosers,
batch.stopping_criterias, batch.stopping_criterias,
batch.all_input_ids, batch.all_input_ids,
) )
...@@ -564,7 +611,6 @@ class FlashCausalLM(Model): ...@@ -564,7 +611,6 @@ class FlashCausalLM(Model):
# For each member of the batch # For each member of the batch
for i, ( for i, (
input_length, input_length,
next_token_chooser,
stopping_criteria, stopping_criteria,
all_input_ids, all_input_ids,
) in enumerate(iterator): ) in enumerate(iterator):
...@@ -573,21 +619,6 @@ class FlashCausalLM(Model): ...@@ -573,21 +619,6 @@ class FlashCausalLM(Model):
end_index = cumulative_length + input_length end_index = cumulative_length + input_length
if prefill: if prefill:
# Prefill mode
# out is of shape [cumulative_sequence_lengths, vocab_size]
# only take last token logit
logits = out[end_index - 1 : end_index]
# Create all_input_ids_tensor that will be used by token warpers (for example, RepetitionPenalty)
all_input_ids_tensor = batch.input_ids.new_empty(
input_length + stopping_criteria.max_new_tokens
)
# Copy from batch.input_ids to all_input_ids_tensor
all_input_ids_tensor[:input_length] = batch.input_ids[
start_index:end_index
]
batch.all_input_ids_tensor.append(all_input_ids_tensor)
# Initialize position_ids # Initialize position_ids
# In decode, we do not need this as we can just increment position ids # In decode, we do not need this as we can just increment position ids
next_position_ids[i] = batch.position_ids[end_index - 1] next_position_ids[i] = batch.position_ids[end_index - 1]
...@@ -603,25 +634,8 @@ class FlashCausalLM(Model): ...@@ -603,25 +634,8 @@ class FlashCausalLM(Model):
prefill_tokens_indices = batch.input_ids[ prefill_tokens_indices = batch.input_ids[
start_index + 1 : end_index start_index + 1 : end_index
] ]
else:
# Decode mode
# out is of shape [batch_size, vocab_size]
logits = out[i].view(1, -1)
all_input_ids_tensor = batch.all_input_ids_tensor[i] batch.all_input_ids_tensor[i, input_length] = next_input_ids[i]
# Select next token
next_token_id, logprob = next_token_chooser(
all_input_ids_tensor[None, :input_length], logits
)
# Add to all_input_ids_tensor
next_token_id_squeezed = next_token_id.view(1)
all_input_ids_tensor[input_length] = next_token_id_squeezed
# Set values
next_input_ids[i] = next_token_id_squeezed
next_token_logprobs[i] = logprob[-1, next_token_id].view(1)
cumulative_length += input_length cumulative_length += input_length
...@@ -651,10 +665,11 @@ class FlashCausalLM(Model): ...@@ -651,10 +665,11 @@ class FlashCausalLM(Model):
batch.input_lengths, batch.input_lengths,
batch.prefix_offsets, batch.prefix_offsets,
batch.read_offsets, batch.read_offsets,
batch.next_token_choosers,
batch.stopping_criterias, batch.stopping_criterias,
batch.all_input_ids, batch.all_input_ids,
batch.all_input_ids_tensor, batch.all_input_ids_tensor,
batch.next_token_chooser.do_sample,
batch.next_token_chooser.seeds,
next_token_ids, next_token_ids,
next_token_logprobs, next_token_logprobs,
) )
...@@ -665,10 +680,11 @@ class FlashCausalLM(Model): ...@@ -665,10 +680,11 @@ class FlashCausalLM(Model):
input_length, input_length,
prefix_offset, prefix_offset,
read_offset, read_offset,
next_token_chooser,
stopping_criteria, stopping_criteria,
all_input_ids, all_input_ids,
all_input_ids_tensor, all_input_ids_tensor,
do_sample,
seed,
next_token_id, next_token_id,
next_token_logprob, next_token_logprob,
) in enumerate(iterator): ) in enumerate(iterator):
...@@ -702,14 +718,11 @@ class FlashCausalLM(Model): ...@@ -702,14 +718,11 @@ class FlashCausalLM(Model):
output_text = self.decode( output_text = self.decode(
all_input_ids[-stopping_criteria.current_tokens :] all_input_ids[-stopping_criteria.current_tokens :]
) )
# Get seed
if isinstance(next_token_chooser.choice, Sampling):
seed = next_token_chooser.choice.seed
else:
seed = None
generated_text = GeneratedText( generated_text = GeneratedText(
output_text, stopping_criteria.current_tokens, reason, seed output_text,
stopping_criteria.current_tokens,
reason,
seed if do_sample else None,
) )
else: else:
generated_text = None generated_text = None
...@@ -751,8 +764,9 @@ class FlashCausalLM(Model): ...@@ -751,8 +764,9 @@ class FlashCausalLM(Model):
batch.prefix_offsets[i] = prefix_offset batch.prefix_offsets[i] = prefix_offset
batch.read_offsets[i] = read_offset batch.read_offsets[i] = read_offset
batch.all_input_ids[i] = all_input_ids batch.all_input_ids[i] = all_input_ids
batch.max_seqlen = batch.max_seqlen + 1
cumulative_length += input_length cumulative_length += input_length
batch.max_seqlen = batch.max_seqlen + 1
# No need to return a batch if we know that all requests stopped # No need to return a batch if we know that all requests stopped
return generations, batch if not stopped else None return generations, batch if not stopped else None
server/text_generation_server/models/galactica.py
View file @ 62f91f78
...@@ -89,6 +89,7 @@ class GalacticaCausalLMBatch(CausalLMBatch): ...@@ -89,6 +89,7 @@ class GalacticaCausalLMBatch(CausalLMBatch):
cls, cls,
pb: generate_pb2.Batch, pb: generate_pb2.Batch,
tokenizer: PreTrainedTokenizerBase, tokenizer: PreTrainedTokenizerBase,
dtype: torch.dtype,
device: torch.device, device: torch.device,
) -> "GalacticaCausalLMBatch": ) -> "GalacticaCausalLMBatch":
inputs = [] inputs = []
......
server/text_generation_server/models/seq2seq_lm.py
View file @ 62f91f78
...@@ -71,6 +71,7 @@ class Seq2SeqLMBatch(Batch): ...@@ -71,6 +71,7 @@ class Seq2SeqLMBatch(Batch):
cls, cls,
pb: generate_pb2.Batch, pb: generate_pb2.Batch,
tokenizer: PreTrainedTokenizerBase, tokenizer: PreTrainedTokenizerBase,
dtype: torch.dtype,
device: torch.device, device: torch.device,
) -> "Seq2SeqLMBatch": ) -> "Seq2SeqLMBatch":
"""Convert a text_generation_server.v1.Batch protobuf to a Seq2SeqLMBatch""" """Convert a text_generation_server.v1.Batch protobuf to a Seq2SeqLMBatch"""
......
server/text_generation_server/models/types.py
View file @ 62f91f78
...@@ -21,6 +21,7 @@ class Batch(ABC): ...@@ -21,6 +21,7 @@ class Batch(ABC):
cls, cls,
pb: generate_pb2.Batch, pb: generate_pb2.Batch,
tokenizer: PreTrainedTokenizerBase, tokenizer: PreTrainedTokenizerBase,
dtype: torch.dtype,
device: torch.device, device: torch.device,
) -> "Batch": ) -> "Batch":
raise NotImplementedError raise NotImplementedError
......
server/text_generation_server/server.py
View file @ 62f91f78
...@@ -55,7 +55,7 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer): ...@@ -55,7 +55,7 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
async def Prefill(self, request, context): async def Prefill(self, request, context):
batch = self.model.batch_type.from_pb( batch = self.model.batch_type.from_pb(
request.batch, self.model.tokenizer, self.model.device request.batch, self.model.tokenizer, self.model.dtype, self.model.device
) )
generations, next_batch = self.model.generate_token(batch) generations, next_batch = self.model.generate_token(batch)
......
server/text_generation_server/utils/__init__.py
View file @ 62f91f78
...@@ -9,12 +9,13 @@ from text_generation_server.utils.hub import ( ...@@ -9,12 +9,13 @@ from text_generation_server.utils.hub import (
RevisionNotFoundError, RevisionNotFoundError,
) )
from text_generation_server.utils.tokens import ( from text_generation_server.utils.tokens import (
Greedy,
NextTokenChooser, NextTokenChooser,
Sampling, HeterogeneousNextTokenChooser,
StoppingCriteria, StoppingCriteria,
StopSequenceCriteria, StopSequenceCriteria,
FinishReason, FinishReason,
Sampling,
Greedy,
) )
__all__ = [ __all__ = [
...@@ -25,6 +26,7 @@ __all__ = [ ...@@ -25,6 +26,7 @@ __all__ = [
"weight_hub_files", "weight_hub_files",
"download_weights", "download_weights",
"EntryNotFoundError", "EntryNotFoundError",
"HeterogeneousNextTokenChooser",
"LocalEntryNotFoundError", "LocalEntryNotFoundError",
"RevisionNotFoundError", "RevisionNotFoundError",
"Greedy", "Greedy",
......
server/text_generation_server/utils/convert.py
View file @ 62f91f78
import concurrent
import time
import datetime import datetime
import torch import torch
from concurrent.futures import ThreadPoolExecutor
from collections import defaultdict from collections import defaultdict
from datetime import timedelta
from loguru import logger from loguru import logger
from pathlib import Path from pathlib import Path
from safetensors.torch import load_file, save_file from safetensors.torch import save_file
from safetensors import safe_open from safetensors import safe_open
from typing import Dict, List from typing import Dict, List
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment