Update modelling_magi.py

modelling_magi.py

@@ -181,7 +181,7 @@ class MagiModel(PreTrainedModel):
 
         return crop_embeddings_for_batch
     
-    def predict_ocr(self, images, crop_bboxes, move_to_device_fn=None, use_tqdm=False, batch_size=32):
+    def predict_ocr(self, images, crop_bboxes, move_to_device_fn=None, use_tqdm=False, batch_size=32, max_new_tokens=64):
         assert not self.config.disable_ocr
         move_to_device_fn = self.move_to_device if move_to_device_fn is None else move_to_device_fn
 
@@ -207,7 +207,7 @@ class MagiModel(PreTrainedModel):
             pbar = range(0, len(crops_per_image), batch_size)
         for i in pbar:
             crops = crops_per_image[i:i+batch_size]
-            generated_ids = self.ocr_model.generate(crops)
+            generated_ids = self.ocr_model.generate(crops, max_new_tokens=max_new_tokens)
             generated_texts = self.processor.postprocess_ocr_tokens(generated_ids)
             all_generated_texts.extend(generated_texts)