Add stop sequence to text generation pipeline (#18444)
This commit is contained in:
@@ -147,6 +147,18 @@ class TextGenerationPipelineTests(unittest.TestCase, metaclass=PipelineTestCaseM
|
||||
text_generator = TextGenerationPipeline(model=model, tokenizer=tokenizer)
|
||||
return text_generator, ["This is a test", "Another test"]
|
||||
|
||||
def test_stop_sequence_stopping_criteria(self):
|
||||
prompt = """Hello I believe in"""
|
||||
text_generator = pipeline("text-generation", model="hf-internal-testing/tiny-random-gpt2")
|
||||
output = text_generator(prompt)
|
||||
self.assertEqual(
|
||||
output,
|
||||
[{"generated_text": "Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe"}],
|
||||
)
|
||||
|
||||
output = text_generator(prompt, stop_sequence=" fe")
|
||||
self.assertEqual(output, [{"generated_text": "Hello I believe in fe"}])
|
||||
|
||||
def run_pipeline_test(self, text_generator, _):
|
||||
model = text_generator.model
|
||||
tokenizer = text_generator.tokenizer
|
||||
|
||||
Reference in New Issue
Block a user