fixess

geniusrise · Feb 15, 2024 · 9bc1d5f · 9bc1d5f
1 parent c448a64
commit 9bc1d5f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/geniusrise_vision/ocr/api.py b/geniusrise_vision/ocr/api.py
@@ -78,7 +78,7 @@ def process_huggingface_models(self, image: Image.Image, use_easyocr_bbox: bool)
         # Convert PIL Image to Tensor
         pixel_values = self.processor(images=image, return_tensors="pt").pixel_values
 
-        if "nougat" in self.model_name.lower():
+        if "nougat" in self.model_name.lower() or "donut" in self.model_name.lower():
             task_prompt = "<s_cord-v2>"
             decoder_input_ids = self.processor.tokenizer(
                 task_prompt, add_special_tokens=False, return_tensors="pt"