Create handler.py

handler.py

@@ -0,0 +1,79 @@
+from copy import deepcopy
+from typing import Any, Dict
+
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers.image_utils import load_image
+
+IMAGE_TOKENS = "<image_start><image><image_end>"
+SEPARATOR = "\n"
+
+
+class EndpointHandler:
+    def __init__(
+        self,
+        model_dir: str = "alvarobartt/Magma-8B",
+        **kwargs: Any,  # type: ignore
+    ) -> None:
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_dir, trust_remote_code=True, torch_dtype=torch.bfloat16
+        ).eval()
+        self.processor = AutoProcessor.from_pretrained(
+            model_dir, trust_remote_code=True
+        )
+
+    def __call__(self, data: Dict[str, Any]) -> Any:
+        if "messages" not in data:
+            raise ValueError(
+                "The request body must contain a key 'messages' with a list of messages."
+            )
+
+        messages, images = [], []
+        for message in data["messages"]:
+            if isinstance(list, message["content"]):
+                new_message = {"role": message["role"], "content": ""}
+                for content in message["content"]:
+                    if content["type"] == "text":
+                        new_message["content"] += content["text"]
+                    elif content["type"] == "image_url":
+                        images.append(load_image(content["image_url"]["url"]))
+                        if new_message["content"].count(
+                            f"{IMAGE_TOKENS}{SEPARATOR}"
+                        ) < len(images):
+                            new_message["content"] = (
+                                f"{IMAGE_TOKENS}{SEPARATOR}" + new_message["content"]
+                            )
+            else:
+                messages.append(
+                    {"role": message["role"], "content": message["content"]}
+                )
+
+        data.pop("messages")
+
+        prompt = self.processor.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+
+        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)
+        inputs = inputs.to("cuda").to(torch.bfloat16)
+
+        generation_args = {
+            "max_new_tokens": 128,
+            "temperature": 0.0,
+            "do_sample": False,
+            "use_cache": True,
+            "num_beams": 1,
+        }
+        generation_args.update(data)
+
+        with torch.inference_mode():
+            generate_ids = self.model.generate(**inputs, **generation_args)
+
+        generate_ids = generate_ids[:, inputs["input_ids"].shape[-1] :]
+        response = self.processor.decode(
+            generate_ids[0], skip_special_tokens=True
+        ).strip()
+
+        return {"generated_text": response}