Update routers/donut_inference.py

routers/donut_inference.py

@@ -7,10 +7,11 @@ from functools import lru_cache
 import os
 
 
-@lru_cache(maxsize=1)
-def load_model():
-    processor = DonutProcessor.from_pretrained(settings.processor)
-    model = VisionEncoderDecoderModel.from_pretrained(settings.model)
+@lru_cache(maxsize=10)
+def load_model(shipper_id: str):
+    temp_settings = Settings(shipper_id=shipper_id)
+    processor = DonutProcessor.from_pretrained(temp_settings.processor)
+    model = VisionEncoderDecoderModel.from_pretrained(temp_settings.model)
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model.to(device)
@@ -18,20 +19,37 @@ def load_model():
     return processor, model, device
 
 
-def process_document_donut(image):
+def process_document_donut(image, shipper_id="default_shipper"):
     worker_pid = os.getpid()
     print(f"Handling inference request with worker PID: {worker_pid}")
 
     start_time = time.time()
 
-    processor, model, device = load_model()
+    processor, model, device = load_model(shipper_id)
 
     # prepare encoder inputs
     pixel_values = processor(image, return_tensors="pt").pixel_values
 
     # prepare decoder inputs
     task_prompt = "<s_cord-v2>"
-    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
+    decoder_input_ids = processor.tokenizer(
+        task_prompt,
+        add_special_tokens=False,
+        return_tensors="pt"
+    ).input_ids
+
+    outputs = model.generate(
+        pixel_values.to(device),
+        decoder_input_ids=decoder_input_ids.to(device),
+        max_length=model.decoder.config.max_position_embeddings,
+        early_stopping=True,
+        pad_token_id=processor.tokenizer.pad_token_id,
+        eos_token_id=processor.tokenizer.eos_token_id,
+        use_cache=True,
+        num_beams=1,
+        bad_words_ids=[[processor.tokenizer.unk_token_id]],
+        return_dict_in_generate=True,
+    )
 
     # generate answer
     outputs = model.generate(