Create handler.py

handler.py

@@ -0,0 +1,92 @@
+# handler.py
+
+from typing import Any, Dict
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_inference_toolkit.logging import logger
+
+
+class EndpointHandler:
+    """
+    Custom handler for Hugging Face Inference Endpoints.
+    """
+
+    def __init__(self, model_dir: str, **kwargs: Any):
+        """
+        Called once at startup. `model_dir` is the path to the downloaded model repository.
+        You can load your model, tokenizer, etc. here.
+        """
+        logger.info(f"Initializing EndpointHandler with model_dir = {model_dir}")
+
+        # load tokenizer and model from the model_dir or hub
+        # we assume model files are present in model_dir
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=False)
+            self.model = AutoModelForCausalLM.from_pretrained(model_dir)
+        except Exception as e:
+            # fallback: try loading from hub name
+            hf_model_id = os.getenv("MODEL_ID") or "Ephraimmm/PIDGIN_gemma-3"
+            logger.warn(f"Loading model from hub id {hf_model_id} because fallback. Exception: {e}")
+            self.tokenizer = AutoTokenizer.from_pretrained(hf_model_id, use_fast=False)
+            self.model = AutoModelForCausalLM.from_pretrained(hf_model_id)
+
+        # put the model on GPU if available
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = self.model.to(self.device)
+        self.model.eval()
+
+        # You could also load extra assets, set up config, etc.
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Called on each API request. `data` is the parsed JSON body from the request.
+        Must return a JSON-serializable structure.
+        The `data` dict always has an “inputs” key (unless you change it).
+        """
+        logger.info(f"EndpointHandler received data: {data}")
+
+        # Extract the prompt (the "inputs" key)
+        prompt = data.get("inputs")
+        if prompt is None:
+            raise ValueError("Request JSON must contain an `inputs` field with the prompt string.")
+
+        # optional: support generation parameters from the payload
+        max_new_tokens = data.get("parameters", {}).get("max_new_tokens", 50)
+        temperature = data.get("parameters", {}).get("temperature", 1.0)
+        top_k = data.get("parameters", {}).get("top_k", 50)
+        top_p = data.get("parameters", {}).get("top_p", 1.0)
+        do_sample = data.get("parameters", {}).get("do_sample", True)
+
+        # Tokenize
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+
+        # Generate
+        output = self.model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            do_sample=do_sample,
+        )
+
+        # Decode
+        full = self.tokenizer.decode(output[0], skip_special_tokens=True)
+
+        # Optionally remove the original prompt from the returned text to send only the continuation
+        continuation = full[len(prompt):] if full.startswith(prompt) else full
+
+        # Prepare response in the shape expected by Inference Endpoints
+        # The “predictions” key is often used (the default inference service expects it)
+        return {
+            "predictions": [full],
+            "continuations": [continuation],
+        }
+
+
+# For local testing (not used in production)
+if __name__ == "__main__":
+    # Quick local test
+    h = EndpointHandler(model_dir=".")
+    resp = h({"inputs": "Once upon a time", "parameters": {"max_new_tokens": 20}})
+    print("Response:", resp)