huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 04:14:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix top_n_tokens > 0

Browse Source 
Fix issue #1340
This commit is contained in:
gduhamel 2024-01-19 20:50:17 +01:00
parent 3ccb3bb0b5
commit fd8b42678d
5 changed files with 7 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
server/text_generation_server/models/causal_lm.py
View File

@ -719,7 +719,7 @@ class CausalLM(Model):
[next_token_id_squeezed.item() in self.all_special_ids],
),
generated_text,
top_tokens,
[top_tokens],
)
generations.append(generation)

2
server/text_generation_server/models/flash_causal_lm.py
View File

@ -1089,7 +1089,7 @@ class FlashCausalLM(Model):
[nid in self.all_special_ids for nid in _next_token_ids],
),
generated_text,
top_tokens,
[top_tokens],
)
generations.append(generation)

2
server/text_generation_server/models/idefics_causal_lm.py
View File

@ -809,7 +809,7 @@ class IdeficsCausalLM(Model):
[next_token_id_squeezed.item() in self.all_special_ids],
),
generated_text,
top_tokens,
[top_tokens],
)
generations.append(generation)

2
server/text_generation_server/models/seq2seq_lm.py
View File

@ -773,7 +773,7 @@ class Seq2SeqLM(Model):
[next_token_id_squeezed.item() in self.all_special_ids],
),
generated_text,
top_tokens,
[top_tokens],
)
generations.append(generation)

4
server/text_generation_server/models/types.py
View File

@ -95,5 +95,7 @@ class Generation:
generated_text=self.generated_text.to_pb()
if self.generated_text is not None
else None,
top_tokens=self.top_tokens.to_pb() if self.top_tokens is not None else None,
top_tokens=([top_token.to_pb() for top_token in self.top_tokens]
if self.top_tokens is not None
else None),
)