Create handler.py

handler.py

@@ -0,0 +1,189 @@
+# handler.py
+# Hugging Face Inference Endpoint custom handler for Mongolian GPT-2 summarization
+# Input JSON:
+#   {
+#     "inputs": "ARTICLE TEXT ...",
+#     "parameters": {
+#        "max_new_tokens": 160,
+#        "num_beams": 4,
+#        "do_sample": false,
+#        "no_repeat_ngram_size": 3,
+#        "length_penalty": 1.0,
+#        "temperature": 1.0,
+#        "top_p": 1.0,
+#        "top_k": 50,
+#        "return_full_text": false
+#     }
+#   }
+# Output JSON:
+#   { "summary_text": "...", "used_new_tokens": 152, "requested_new_tokens": 160 }
+
+from typing import Any, Dict, List, Union
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Mongolian instruction + prompt template used during training
+INSTRUCTION = "Дараах бичвэрийг хураангуйлж бич."
+PROMPT_TEMPLATE = (
+    "### Даалгавар:\n"
+    f"{INSTRUCTION}\n\n"
+    "### Бичвэр:\n{article}\n\n"
+    "### Хураангуй:\n"
+)
+
+def _select_dtype() -> torch.dtype:
+    if torch.cuda.is_available():
+        # Prefer bf16 if supported; otherwise use fp16
+        return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    return torch.float32
+
+class EndpointHandler:
+    """
+    Custom handler for HF Inference Endpoints:
+    - __init__(path): loads model assets from `path`
+    - __call__(data): performs generation given {"inputs": ..., "parameters": {...}}
+    """
+    def __init__(self, path: str = ""):
+        # Device & dtype
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.dtype = _select_dtype()
+
+        # Load tokenizer/model from the repository directory
+        self.tokenizer = AutoTokenizer.from_pretrained(path, use_fast=True)
+        # Decoder-only model requires left padding for correct generation
+        self.tokenizer.padding_side = "left"
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        self.model = AutoModelForCausalLM.from_pretrained(
+            path,
+            torch_dtype=self.dtype,
+        ).to(self.device)
+        # Safer attention path on many endpoint stacks
+        self.model.config.attn_implementation = "eager"
+        self.model.config.pad_token_id = self.tokenizer.pad_token_id
+        self.model.config.eos_token_id = self.tokenizer.eos_token_id
+        self.model.eval()
+
+        # Read max context from config (GPT-2 default is 1024)
+        self.max_context = getattr(self.model.config, "max_position_embeddings", 1024)
+
+    def _build_prompt(self, article: str) -> str:
+        return PROMPT_TEMPLATE.format(article=article.strip())
+
+    def _prepare_inputs(
+        self,
+        articles: List[str],
+        requested_new: int
+    ):
+        """
+        Tokenize prompts so that prompt_len + max_new_tokens <= max_context.
+        We first clamp requested_new, then tokenize with truncation=max_context - requested_new.
+        """
+        # Basic safety clamps
+        requested_new = int(max(1, min(requested_new, 512)))
+        max_len_for_prompt = max(1, self.max_context - requested_new)
+
+        prompts = [self._build_prompt(a) for a in articles]
+        enc = self.tokenizer(
+            prompts,
+            add_special_tokens=False,
+            truncation=True,
+            max_length=max_len_for_prompt,
+            return_tensors="pt",
+            padding=True,  # uses left padding because tokenizer.padding_side="left"
+        )
+        enc = {k: v.to(self.device) for k, v in enc.items()}
+
+        # Compute per-example available space and adjust new tokens if needed
+        input_lens = enc["attention_mask"].sum(dim=1).tolist()
+        per_example_new = []
+        for L in input_lens:
+            available = max(0, self.max_context - int(L))
+            per_example_new.append(max(1, min(requested_new, available)))
+
+        return enc, per_example_new, prompts
+
+    @torch.no_grad()
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        # Accept either {"inputs": "..."} or {"inputs": ["...", "..."]}
+        raw_inputs: Union[str, List[str], Dict[str, Any]] = data.get("inputs", "")
+        params: Dict[str, Any] = data.get("parameters", {}) or {}
+
+        # Default generation hyperparameters (aligned with training)
+        req_new = int(params.get("max_new_tokens", 160))
+        num_beams = int(params.get("num_beams", 4))
+        do_sample = bool(params.get("do_sample", False))
+        no_repeat = int(params.get("no_repeat_ngram_size", 3))
+        length_penalty = float(params.get("length_penalty", 1.0))
+        temperature = float(params.get("temperature", 1.0))
+        top_p = float(params.get("top_p", 1.0))
+        top_k = int(params.get("top_k", 50))
+        return_full_text = bool(params.get("return_full_text", False))
+
+        # Normalize inputs to a list of strings
+        if isinstance(raw_inputs, str):
+            articles = [raw_inputs]
+        elif isinstance(raw_inputs, list):
+            if not all(isinstance(x, str) for x in raw_inputs):
+                raise ValueError("All elements of 'inputs' must be strings.")
+            articles = raw_inputs
+        else:
+            # Accept {"article": "..."} as a courtesy
+            maybe_article = data.get("article")
+            if isinstance(maybe_article, str):
+                articles = [maybe_article]
+            else:
+                raise ValueError("Expect 'inputs' as a string or list of strings.")
+
+        # Tokenize prompts and cap new tokens per example
+        enc, per_example_new, prompts = self._prepare_inputs(articles, req_new)
+
+        # Generate (batched)
+        gen_out = self.model.generate(
+            **enc,
+            max_new_tokens=max(per_example_new),  # upper bound; actual stopping still respects EOS
+            num_beams=num_beams,
+            do_sample=do_sample,
+            no_repeat_ngram_size=no_repeat,
+            length_penalty=length_penalty,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            pad_token_id=self.tokenizer.pad_token_id,
+            eos_token_id=self.tokenizer.eos_token_id,
+            early_stopping=True,
+        )
+
+        # Decode and postprocess per-item (cut after the prompt if needed)
+        decoded = self.tokenizer.batch_decode(gen_out, skip_special_tokens=True)
+
+        results = []
+        for i, text in enumerate(decoded):
+            if return_full_text:
+                full = text.strip()
+                # Try to extract summary part for convenience too
+                split_key = "### Хураангуй:\n"
+                summary = full.split(split_key, 1)[-1].strip() if split_key in full else full
+            else:
+                # Remove the prompt prefix, return only the generated summary
+                prefix = prompts[i]
+                if text.startswith(prefix):
+                    summary = text[len(prefix):].strip()
+                else:
+                    # Fallback split on the marker
+                    split_key = "### Хураангуй:\n"
+                    summary = text.split(split_key, 1)[-1].strip() if split_key in text else text.strip()
+                full = None
+
+            results.append({
+                "summary_text": summary,
+                "used_new_tokens": per_example_new[i],
+                "requested_new_tokens": req_new,
+                **({"full_text": full} if return_full_text else {})
+            })
+
+        # If the input was a single string, return a single object
+        if isinstance(raw_inputs, str):
+            return results[0]
+        return {"results": results}