handler from tut

handler.py

@@ -1,168 +1,25 @@
-from typing import Dict, List, Any, Optional, Union
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor
-import torch
-import base64
-from io import BytesIO
-from PIL import Image
-import requests
+from typing import Any, Dict
 
 
 class EndpointHandler:
-    def __init__(self, path=""):
-        # If path is empty, use the GSAI-ML/LLaDA-8B-Instruct model
-        if not path:
-            path = "GSAI-ML/LLaDA-8B-Instruct"
-
-        print(f"Loading model from {path}...")
-
-        # Load model with half precision to save memory
-        self.model = AutoModelForCausalLM.from_pretrained(
-            path, torch_dtype=torch.float16, device_map="auto"
+    def __init__(self, model_dir: str, **kwargs: Any) -> None:
+			self.model = AutoModel.from_pretrained(
+            model_dir,
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            use_flash_attn=False,
+            trust_remote_code=True,
+            device_map=split_model(),
+        ).eval()
+
+				self.tokenizer = AutoTokenizer.from_pretrained(
+            model_dir, trust_remote_code=True, use_fast=False
         )
 
-        # Load tokenizer
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
-
-        # Load processor for handling images
-        self.processor = AutoProcessor.from_pretrained(path)
-
-        # Ensure pad token is properly set
-        if self.tokenizer.pad_token_id is None:
-            if (
-                hasattr(self.tokenizer, "eos_token_id")
-                and self.tokenizer.eos_token_id is not None
-            ):
-                self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
-                self.tokenizer.pad_token = self.tokenizer.eos_token
-            else:
-                # Fallback to a common pad token
-                self.tokenizer.pad_token_id = 0
-                self.tokenizer.pad_token = self.tokenizer.convert_ids_to_tokens(0)
-
-        print(f"Model loaded successfully. Pad token ID: {self.tokenizer.pad_token_id}")
-
-    def _load_image(self, image_data: Union[str, bytes]) -> Image.Image:
-        """Load image from URL or base64 encoded string"""
-        if isinstance(image_data, str):
-            if image_data.startswith("http"):
-                # Load from URL
-                response = requests.get(image_data, stream=True)
-                response.raise_for_status()
-                return Image.open(BytesIO(response.content))
-            elif image_data.startswith("data:image"):
-                # Handle base64 encoded image
-                base64_data = image_data.split(",")[1]
-                image_bytes = base64.b64decode(base64_data)
-                return Image.open(BytesIO(image_bytes))
-            else:
-                # Assume it's a base64 string without the prefix
-                try:
-                    image_bytes = base64.b64decode(image_data)
-                    return Image.open(BytesIO(image_bytes))
-                except Exception as e:
-                    raise ValueError(f"Invalid image data format: {e}")
-        elif isinstance(image_data, bytes):
-            return Image.open(BytesIO(image_data))
-        else:
-            raise ValueError(f"Unsupported image data type: {type(image_data)}")
-
-    def _format_prompt(self, text: str, system_prompt: Optional[str] = None) -> str:
-        """Format the prompt according to LLaDA's expected format"""
-        # Default system prompt for LLaDA if none provided
-        if system_prompt is None:
-            system_prompt = (
-                "You are a helpful AI assistant that can understand images and text."
-            )
-
-        # Format the prompt following LLaDA's expected structure
-        formatted_prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
-        return formatted_prompt
-
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, List[Any]]:
-        """Handle inference requests"""
-        # Extract inputs and parameters from request data
-        inputs = data.pop("inputs", data)
-        parameters = data.pop("parameters", {})
-
-        # Extract image data if present
-        image_data = parameters.get("image", None)
-        system_prompt = parameters.get("system_prompt", None)
-
-        # Extract generation parameters with sensible defaults
-        max_new_tokens = parameters.get("max_new_tokens", 256)
-        temperature = parameters.get("temperature", 0.7)
-        top_p = parameters.get("top_p", 0.95)
-        do_sample = parameters.get("do_sample", True)
-
-        # Convert single string input to list for consistent handling
-        if isinstance(inputs, str):
-            inputs = [inputs]
-
-        # Process each input
-        generated_texts = []
-        for input_text in inputs:
-            # Format the prompt according to LLaDA's expected format
-            formatted_prompt = self._format_prompt(input_text, system_prompt)
-
-            if image_data:
-                try:
-                    # Process image if present
-                    image = self._load_image(image_data)
-                    inputs_processor = self.processor(
-                        text=formatted_prompt, images=image, return_tensors="pt"
-                    )
-
-                    # Move inputs to the same device as the model
-                    for k, v in inputs_processor.items():
-                        if isinstance(v, torch.Tensor):
-                            inputs_processor[k] = v.to(self.model.device)
-
-                    # Generate text with image context
-                    with torch.no_grad():
-                        outputs = self.model.generate(
-                            **inputs_processor,
-                            max_new_tokens=max_new_tokens,
-                            temperature=temperature,
-                            top_p=top_p,
-                            do_sample=do_sample,
-                            pad_token_id=self.tokenizer.pad_token_id,
-                        )
-
-                    # Decode generated text
-                    generated_text = self.tokenizer.decode(
-                        outputs[0], skip_special_tokens=True
-                    )
-                    generated_texts.append(generated_text)
-
-                except Exception as e:
-                    # If image processing fails, fall back to text-only
-                    print(
-                        f"Error processing image: {e}. Falling back to text-only processing."
-                    )
-                    image_data = None
-
-            if not image_data:
-                # Text-only processing
-                input_tokens = self.tokenizer(formatted_prompt, return_tensors="pt").to(
-                    self.model.device
-                )
-
-                # Generate text
-                with torch.no_grad():
-                    outputs = self.model.generate(
-                        **input_tokens,
-                        max_new_tokens=max_new_tokens,
-                        temperature=temperature,
-                        top_p=top_p,
-                        do_sample=do_sample,
-                        pad_token_id=self.tokenizer.pad_token_id,
-                    )
+    def __call__(self, data: Dict[str, Any]) -> Any:
+			logger.info(f"Received incoming request with {data=}")
 
-                # Decode generated text
-                generated_text = self.tokenizer.decode(
-                    outputs[0], skip_special_tokens=True
-                )
-                generated_texts.append(generated_text)
 
-        # Return results in expected format
-        return {"generated_text": generated_texts}
+if __name__ == "__main__":
+    handler = EndpointHandler(model_dir="GSAI-ML/LLaDA-8B-Instruct")
+    print(handler)