try fix gen

infer.py

@@ -203,13 +203,14 @@ class TikzGenerator:
             top_p=top_p,
             top_k=top_k,
             num_return_sequences=1,
-            max_length=self.pipeline.tokenizer.model_max_length, # type: ignore
+            # max_length=self.pipeline.tokenizer.model_max_length, # type: ignore
             do_sample=True,
             return_full_text=False,
             streamer=TextStreamer(self.pipeline.tokenizer, # type: ignore
                 skip_prompt=True,
                 skip_special_tokens=True
             ),
+            max_new_tokens=1024,
         )
 
         if not stream:
@@ -218,8 +219,11 @@ class TikzGenerator:
     def generate(self, image: Image.Image, **generate_kwargs):
         prompt = "Assistant helps to write down the TikZ code for the user's image. USER: <image>\nWrite down the TikZ code to draw the diagram shown in the lol. ASSISTANT:"
         tokenizer = self.pipeline.tokenizer
+        print('starting generation')
         text = self.pipeline(image, prompt=prompt, generate_kwargs=(self.default_kwargs | generate_kwargs))[0]["generated_text"] # type: ignore
 
+        print('text generated: ', text)  # TODO: remove
+
         if self.clean_up_output:
             for token in reversed(tokenizer.tokenize(prompt)): # type: ignore
                 # remove leading characters because skip_special_tokens in pipeline
@@ -236,7 +240,9 @@ class TikzGenerator:
             for artifact, replacement in artifacts.items():
                 text = sub(artifact, replacement, text) # type: ignore
 
-        return text
+        print('cleaned text: ', text)
+
+        return TikzDocument(text)
 
 
     def __call__(self, *args, **kwargs):