Add logging

Pipfile

@@ -7,6 +7,7 @@ name = "pypi"
 transformers = "*"
 pillow = "*"
 torch = "*"
+logging = "*"
 
 [dev-packages]
 

Pipfile.lock

@@ -1,7 +1,7 @@
 {
     "_meta": {
         "hash": {
-            "sha256": "7a2624bd76968ab76a35343f263a0650cb105c90def7d0a508d9874eaa81f979"
+            "sha256": "d2cc81eabeb4001a0933c7fe68bde6dae34d241899c017a7e1fada250b02d606"
         },
         "pipfile-spec": 6,
         "requires": {
@@ -160,6 +160,13 @@
             "markers": "python_version >= '3.7'",
             "version": "==3.1.4"
         },
+        "logging": {
+            "hashes": [
+                "sha256:26f6b50773f085042d301085bd1bf5d9f3735704db9f37c1ce6d8b85c38f2417"
+            ],
+            "index": "pypi",
+            "version": "==0.4.9.6"
+        },
         "markupsafe": {
             "hashes": [
                 "sha256:00e046b6dd71aa03a41079792f8473dc494d564611a8f89bbbd7cb93295ebdcf",

handler.py

@@ -1,16 +1,19 @@
 from typing import Any, Dict, List
 from transformers import Idefics2Processor, Idefics2ForConditionalGeneration
 import torch
+import logging
 
 
 class EndpointHandler:
     def __init__(self, path=""):
         # Preload all the elements you are going to need at inference.
+        self.logger = logging.getLogger()
+        self.logger.addHandler(logging.StreamHandler())
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.processor = Idefics2Processor.from_pretrained(path)
         self.model = Idefics2ForConditionalGeneration.from_pretrained(path)
         self.model.to(self.device)
-        print("Initialisation finished!")
+        self.logger.info("Initialisation finished!")
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -20,20 +23,25 @@ class EndpointHandler:
         Return:
               A :obj:`list` | `dict`: will be serialized and returned
         """
-        image = data.pop("inputs", data)
-        print("image reached")
+        checkpoints = ""
 
-        # process image
-        inputs = self.processor(images=image, return_tensors="pt").to(self.device)
-        print("inputs reached")
-        generated_ids = self.model.generate(**inputs)
-        print("generated")
+        try:
+            image = data.pop("inputs", data)
+            checkpoints += "image reached\n"
 
-        # run prediction
-        generated_text = self.processor.batch_decode(
-            generated_ids, skip_special_tokens=True
-        )
-        print("decoded")
+            # process image
+            inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+            checkpoints += "inputs reached\n"
+            generated_ids = self.model.generate(**inputs, max_new_tokens=20)
+            checkpoints += "generated\n"
+
+            # run prediction
+            generated_text: List[str] = self.processor.batch_decode(
+                generated_ids, skip_special_tokens=True
+            )
+            checkpoints += "decoded\n"
+        except Exception as e:
+            checkpoints += f"{e}\n"
 
         # decode output
-        return generated_text
+        return generated_text.append(checkpoints)

requirements.txt

@@ -6,6 +6,7 @@ fsspec==2024.6.0; python_version >= '3.8'
 huggingface-hub==0.23.3; python_full_version >= '3.8.0'
 idna==3.7; python_version >= '3.5'
 jinja2==3.1.4; python_version >= '3.7'
+logging==0.4.9.6
 markupsafe==2.1.5; python_version >= '3.7'
 mpmath==1.3.0
 networkx==3.3; python_version >= '3.10'