Merge pull request #168 from tharapalanivel/stop_sequences_fix

🐛 stop_sequences as list of str
caikit · Sep 5, 2023 · d0c18d7 · d0c18d7
2 parents 841f446 + 5f93677
commit d0c18d7
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 11 deletions.
diff --git a/caikit_nlp/modules/text_generation/peft_prompt_tuning.py b/caikit_nlp/modules/text_generation/peft_prompt_tuning.py
@@ -186,7 +186,7 @@ def run(
         exponential_decay_length_penalty: Optional[
             Union[Tuple[int, float], ExponentialDecayLengthPenalty]
         ] = None,
-        stop_sequences: Optional[str] = None,
+        stop_sequences: Optional[List[str]] = None,
     ) -> GeneratedTextResult:
         """
             Run the full text generation model.
@@ -246,7 +246,7 @@ def run_stream_out(
         exponential_decay_length_penalty: Optional[
             Union[Tuple[int, float], ExponentialDecayLengthPenalty]
         ] = None,
-        stop_sequences: Optional[str] = None,
+        stop_sequences: Optional[List[str]] = None,
     ) -> Iterable[GeneratedTextStreamResult]:
         """Run the text generation model with output streaming
 

diff --git a/caikit_nlp/modules/text_generation/peft_tgis_remote.py b/caikit_nlp/modules/text_generation/peft_tgis_remote.py
@@ -181,7 +181,7 @@ def run(
         exponential_decay_length_penalty: Optional[
             Union[Tuple[int, float], ExponentialDecayLengthPenalty]
         ] = None,
-        stop_sequences: Optional[str] = None,
+        stop_sequences: Optional[List[str]] = None,
     ) -> GeneratedTextResult:
         """Run inference against the model running in TGIS.
 
@@ -237,7 +237,7 @@ def run_stream_out(
         exponential_decay_length_penalty: Optional[
             Union[Tuple[int, float], ExponentialDecayLengthPenalty]
         ] = None,
-        stop_sequences: Optional[str] = None,
+        stop_sequences: Optional[List[str]] = None,
     ) -> Iterable[GeneratedTextStreamResult]:
         """Run output stream inferencing against the model running in TGIS
 

diff --git a/caikit_nlp/modules/text_generation/text_generation_tgis.py b/caikit_nlp/modules/text_generation/text_generation_tgis.py
@@ -14,7 +14,7 @@
 
 
 # Standard
-from typing import Iterable, Optional, Tuple, Union
+from typing import Iterable, List, Optional, Tuple, Union
 import os
 
 # First Party
@@ -217,7 +217,7 @@ def run(
         exponential_decay_length_penalty: Optional[
             Union[Tuple[int, float], ExponentialDecayLengthPenalty]
         ] = None,
-        stop_sequences: Optional[str] = None,
+        stop_sequences: Optional[List[str]] = None,
     ) -> GeneratedTextResult:
         """Run inference against the model running in TGIS.
 
@@ -268,7 +268,7 @@ def run_stream_out(
         exponential_decay_length_penalty: Optional[
             Union[Tuple[int, float], ExponentialDecayLengthPenalty]
         ] = None,
-        stop_sequences: Optional[str] = None,
+        stop_sequences: Optional[List[str]] = None,
     ) -> Iterable[GeneratedTextStreamResult]:
         """Run output stream inferencing for text generation module.
 

diff --git a/caikit_nlp/toolkit/text_generation/model_run_utils.py b/caikit_nlp/toolkit/text_generation/model_run_utils.py
@@ -15,7 +15,7 @@
 """Utility functions used for executing run function for text_generation"""
 
 # Standard
-from typing import Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 # Third Party
 from transformers import StoppingCriteria, TextStreamer
@@ -96,7 +96,7 @@
         consist of: (start_index, decay_factor) where start_index
         indicates where penalty starts and decay_factor represents the factor
         of exponential decay
-    stop_sequences: List[str]:
+    stop_sequences: List[str]
         List of strings to be used as stopping criteria
 """
 
@@ -147,7 +147,7 @@ def generate_text_func(
     exponential_decay_length_penalty: Optional[
         Union[Tuple[int, float], ExponentialDecayLengthPenalty]
     ] = None,
-    stop_sequences: Optional[str] = None,
+    stop_sequences: Optional[List[str]] = None,
     **kwargs,
 ):
     """
@@ -255,7 +255,7 @@ def generate_text_func_stream(
     exponential_decay_length_penalty: Optional[
         Union[Tuple[int, float], ExponentialDecayLengthPenalty]
     ] = None,
-    stop_sequences: Optional[str] = None,
+    stop_sequences: Optional[List[str]] = None,
     **kwargs,
 ):
     """