1  """
 2  Generator module tests
 3  """
 4  
 5  import unittest
 6  
 7  from txtai.pipeline import Generator
 8  
 9  
10  class TestGenerator(unittest.TestCase):
11      """
12      Sequences tests.
13      """
14  
15      def testGeneration(self):
16          """
17          Test text pipeline generation
18          """
19  
20          model = Generator("hf-internal-testing/tiny-random-gpt2")
21          start = "Hello, how are"
22  
23          # Test that text is generated
24          self.assertIsNotNone(model(start))