model_run_utils: add/fix type hints

Signed-off-by: Daniele Trifirò <[email protected]>
caikit · Oct 2, 2023 · eab5fcf · eab5fcf
1 parent ef283e4
commit eab5fcf
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/caikit_nlp/toolkit/text_generation/model_run_utils.py b/caikit_nlp/toolkit/text_generation/model_run_utils.py
@@ -19,7 +19,7 @@
 
 # Third Party
 from peft.peft_model import PeftModel
-from transformers import StoppingCriteria, TextStreamer
+from transformers import AutoModel, AutoTokenizer, StoppingCriteria, TextStreamer
 import numpy as np
 import torch
 
@@ -130,10 +130,10 @@ def __iter__(self):
 
 
 def generate_text_func(
-    model,
-    tokenizer,
+    model: "Union[PeftModel, AutoModel]",
+    tokenizer: "AutoTokenizer",
     producer_id: ProducerId,
-    eos_token: str,
+    eos_token: Optional[str],
     text: str,
     max_new_tokens: Optional[int] = 20,
     min_new_tokens: Optional[int] = 0,