Fix formatting of CLI generate command input (#402)

* Fix formatting of CLI generate input

* Update unit test to verify the generate input prompt

* Test doc string

Author: dmsuehir

examples/README.md

@@ -6,6 +6,7 @@ Here are some examples using no-code bash commands:
 * [Text Classification Intel® Transfer Learning Tool CLI Example](cli/text_classification.md)
 * [Image Classification Intel® Transfer Learning Tool CLI Example](cli/image_classification.md)
 * [Vision Anomaly Detection Intel® Transfer Learning Tool CLI Example](cli/image_anomaly_detection.md)
+* [Text Generation Intel® Transfer Learning Tool CLI Example](cli/text_generation.md)
 
 Here are Jupyter notebook examples using low-code Python\* API calls:
 

tests/tools/cli/test_generate_cli.py

@@ -58,18 +58,31 @@ def teardown_class(cls):
                 print("Deleting test directory:", dir)
                 shutil.rmtree(dir)
 
+    @patch("tlt.models.text_generation.pytorch_hf_text_generation_model.PyTorchHFTextGenerationModel.generate")
     @pytest.mark.parametrize('model_name,prompt',
                              [['distilgpt2', 'The size of an apple is'],
-                              ['distilgpt2', 'A large fruit is']])
-    def test_base_generation(self, model_name, prompt):
+                              ['distilgpt2', 'A large fruit is'],
+                              ['distilgpt2',
+                               'The input describes a task.\\n\\nInstruction:\nWrite a song.\\n\\n### Response:\n']])
+    def test_base_generation(self, mock_generate, model_name, prompt):
         """
-        Tests the full workflow for PYT text generation using a custom dataset
+        Tests the CLI generate command for PYT text generation using a HF pretrained model
         """
         runner = CliRunner()
 
+        # Define a dummy response
+        mock_generate.return_value = [prompt + ' so good.']
+
         # Generate a text completion
         result = runner.invoke(generate,
                                ["--model-name", model_name, "--prompt", prompt])
+
+        # Verify that the TLT generate method was called with a properly formatted prompt string
+        assert len(mock_generate.call_args_list) == 1
+        prompt_arg = mock_generate.call_args_list[0][0]
+        assert "\\n" not in prompt_arg
+
+        # Verify that we didn't get any errors
         assert result is not None
         assert result.exit_code == 0
 

tlt/tools/cli/commands/generate.py

@@ -95,10 +95,7 @@ def generate(model_dir, model_name, prompt, temperature, top_p, top_k, repetitio
         if os.path.exists(model_dir):
             model.load_from_directory(model_dir)
 
-    print()
-    print("Prompt:", prompt)
-    print()
-
+    prompt = prompt.replace("\\n", "\n")
     output = model.generate(prompt, temperature=temperature, repetition_penalty=repetition_penalty, top_p=top_p,
                             top_k=top_k, num_beams=num_beams, max_new_tokens=max_new_tokens)
-    print(output)
+    print(*output, sep='\n')