logging

inference.py

@@ -1,4 +1,5 @@
 from transformers import ViltProcessor, ViltForQuestionAnswering, Pix2StructProcessor, Pix2StructForConditionalGeneration, Blip2Processor, Blip2ForConditionalGeneration
+from transformers.utils import logging
 
 class Inference:
     def __init__(self):
@@ -9,12 +10,15 @@ class Inference:
         self.deplot_model = Pix2StructForConditionalGeneration.from_pretrained('google/deplot')
 
     def inference(self, selected, image, text):
+        logger.info(f"selected model {selected}")
         if selected == "Model 1":
             return self.__inference_deplot(image, text)
         elif selected == "Model 2":
             return self.__inference_deplot(image, text)
         elif selected == "Model 3":
             return self.__inference_vilt(image, text)
+        else:
+            logger.warning("Please select a model to make the inference..")
 
     def __inference_vilt(self, image, text):
         encoding = self.vilt_processor(image, text, return_tensors="pt")