Merge pull request #247 from dtrifiro/fix-uncaught-exception

fix uncaught exception in http client

Author: evaline-ju

caikit_nlp/toolkit/text_generation/model_run_utils.py

@@ -27,6 +27,7 @@
 from caikit.core.data_model.producer import ProducerId
 from caikit.core.exceptions import error_handler
 from caikit.interfaces.nlp.data_model import (
+    FinishReason,
     GeneratedTextResult,
     GeneratedTextStreamResult,
     TokenStreamDetails,
@@ -237,16 +238,16 @@ def generate_text_func(
     if (eos_token and tokenizer.decode(generate_ids[0, -1].item()) == eos_token) or (
         generate_ids[0, -1] == tokenizer.eos_token_id
     ):
-        finish_reason = "EOS_TOKEN"
+        finish_reason = FinishReason.EOS_TOKEN
     elif ("stopping_criteria" in gen_optional_params) and (
         gen_optional_params["stopping_criteria"](
             generate_ids,
             None,  # scores, unused by SequenceStoppingCriteria
         )
     ):
-        finish_reason = "STOP_SEQUENCE"
+        finish_reason = FinishReason.STOP_SEQUENCE
     else:
-        finish_reason = "MAX_TOKENS"
+        finish_reason = FinishReason.MAX_TOKENS
 
     return GeneratedTextResult(
         generated_tokens=token_count,

tests/toolkit/text_generation/test_model_run_utils.py

@@ -0,0 +1,45 @@
+# Third Party
+import pytest
+
+# First Party
+from caikit.core.data_model.producer import ProducerId
+from caikit.interfaces.nlp.data_model import GeneratedTextResult
+
+# Local
+from caikit_nlp.toolkit.text_generation.model_run_utils import generate_text_func
+from tests.fixtures import (
+    causal_lm_dummy_model,
+    causal_lm_train_kwargs,
+    seq2seq_lm_dummy_model,
+    seq2seq_lm_train_kwargs,
+)
+
+
+@pytest.mark.parametrize(
+    "model_fixture", ["seq2seq_lm_dummy_model", "causal_lm_dummy_model"]
+)
+@pytest.mark.parametrize(
+    "serialization_method,expected_type",
+    [
+        ("to_dict", dict),
+        ("to_json", str),
+        ("to_proto", GeneratedTextResult._proto_class),
+    ],
+)
+def test_generate_text_func_serialization_json(
+    request,
+    model_fixture,
+    serialization_method,
+    expected_type,
+):
+    model = request.getfixturevalue(model_fixture)
+    generated_text = generate_text_func(
+        model=model.model,
+        tokenizer=model.tokenizer,
+        producer_id=ProducerId("TextGeneration", "0.1.0"),
+        eos_token="<\n>",
+        text="What is the boiling point of liquid Nitrogen?",
+    )
+
+    serialized = getattr(generated_text, serialization_method)()
+    assert isinstance(serialized, expected_type)