Update handler.py

handler.py

@@ -1,7 +1,7 @@
-from copy import deepcopy
 from typing import Any, Dict
 
 import torch
+from huggingface_inference_toolkit.logging import logger
 from transformers import AutoModelForCausalLM, AutoProcessor
 from transformers.image_utils import load_image
 
@@ -23,20 +23,30 @@ class EndpointHandler:
         )
 
     def __call__(self, data: Dict[str, Any]) -> Any:
-        if "messages" not in data:
+        logger.info(f"Received payload with {data}")
+        if "inputs" not in data:
             raise ValueError(
-                "The request body must contain a key 'messages' with a list of messages."
+                "The request body must contain a key 'inputs' with a list of messages."
             )
 
+        logger.info("Processing the messages...")
         messages, images = [], []
-        for message in data["messages"]:
+        for message in data["inputs"]:
+            logger.info(f"Processing {message=}...")
             if isinstance(message["content"], list):
                 new_message = {"role": message["role"], "content": ""}
                 for content in message["content"]:
+                    logger.info(f"{message=} is of type {content['type']}")
                     if content["type"] == "text":
                         new_message["content"] += content["text"]
                     elif content["type"] == "image_url":
                         images.append(load_image(content["image_url"]["url"]))
+                        logger.info(
+                            "Loaded image using `transformers.image_utils.load_image`"
+                        )
+                        logger.info(
+                            "Current {new_message['content']} text if any contains {new_message['content'].count(IMAGE_TOKENS)} image tokens"
+                        )
                         if new_message["content"].count(
                             f"{IMAGE_TOKENS}{SEPARATOR}"
                         ) < len(images):
@@ -48,12 +58,14 @@ class EndpointHandler:
                     {"role": message["role"], "content": message["content"]}
                 )
 
-        data.pop("messages")
+        data.pop("inputs")
 
+        logger.info(f"Applying chat template to {messages=}")
         prompt = self.processor.tokenizer.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=True
         )
 
+        logger.info(f"Processing {len(images)} images...")
         inputs = self.processor(images=images, texts=prompt, return_tensors="pt")
         inputs["pixel_values"] = inputs["pixel_values"].unsqueeze(0)
         inputs["image_sizes"] = inputs["image_sizes"].unsqueeze(0)
@@ -67,13 +79,16 @@ class EndpointHandler:
             "num_beams": 1,
         }
         generation_args.update(data)
+        logger.info(f"Running text generation with the following {generation_args=}")
 
         with torch.inference_mode():
             generate_ids = self.model.generate(**inputs, **generation_args)
 
+        logger.info(f"Generated {generate_ids=}")
         generate_ids = generate_ids[:, inputs["input_ids"].shape[-1] :]
         response = self.processor.decode(
             generate_ids[0], skip_special_tokens=True
         ).strip()
+        logger.info(f"Generated the {response=}")
 
         return {"generated_text": response}