Add Type Argument To Greedy Decoding

apaniukov · Jan 30, 2024 · 7f601d0 · 7f601d0
1 parent 2e0bd47
commit 7f601d0
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/modules/custom_operations/user_ie_extensions/tokenizer/python/openvino_tokenizers/utils.py b/modules/custom_operations/user_ie_extensions/tokenizer/python/openvino_tokenizers/utils.py
@@ -86,9 +86,12 @@ def greedy_decoder(input) -> Model:
     return token_ids.output(0)
 
 
-def add_greedy_decoding(text_generation_model: Model, logits_output: str = LOGITS_OUTPUT_NAME) -> Model:
+def add_greedy_decoding(
+        text_generation_model: Model, logits_output: str = LOGITS_OUTPUT_NAME, output_type: Type = Type.i64
+) -> Model:
     ppp = PrePostProcessor(text_generation_model)
     ppp.output(logits_output).postprocess().custom(greedy_decoder)
+    ppp.output(logits_output).tensor().set_element_type(output_type)
     model = ppp.build()
     model.output(logits_output).tensor.set_names({TOKEN_IDS_OUTPUT_NAME})
     return model