| """ |
| Custom Inference Endpoint handler for Gemma4 text (drive-thru merged model). |
| Requires a recent transformers build that knows `gemma4_text` / Gemma4ForCausalLM. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from typing import Any, Dict, List |
|
|
| import torch |
| from transformers import AutoTokenizer, Gemma4ForCausalLM |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path: str = ""): |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" |
| dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32 |
| self.tokenizer = AutoTokenizer.from_pretrained(path) |
| if self.tokenizer.pad_token is None: |
| self.tokenizer.pad_token = self.tokenizer.eos_token |
| self.model = Gemma4ForCausalLM.from_pretrained( |
| path, |
| torch_dtype=dtype, |
| device_map="auto" if self.device == "cuda" else None, |
| ) |
| if self.device == "cpu": |
| self.model = self.model.to(self.device) |
| self.model.eval() |
|
|
| def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]: |
| |
| if "inputs" in data: |
| prompt = data["inputs"] |
| params = data.get("parameters") or {} |
| elif "messages" in data: |
| messages = data["messages"] |
| system = "" |
| user = "" |
| for m in messages: |
| if m.get("role") == "system": |
| system = m.get("content") or "" |
| elif m.get("role") == "user": |
| user = m.get("content") or "" |
| prompt = ( |
| f"<bos><start_of_turn>system\n{system}<end_of_turn>\n" |
| f"<start_of_turn>user\n{user}<end_of_turn>\n" |
| f"<start_of_turn>model\n" |
| ) |
| params = { |
| "max_new_tokens": data.get("max_tokens") or data.get("max_new_tokens") or 120, |
| "temperature": data.get("temperature", 0.2), |
| } |
| else: |
| prompt = str(data) |
| params = {} |
|
|
| max_new = int(params.get("max_new_tokens") or params.get("max_tokens") or 120) |
| temperature = float(params.get("temperature", 0.2)) |
| do_sample = temperature > 0 |
|
|
| inputs = self.tokenizer(prompt, return_tensors="pt") |
| inputs = {k: v.to(self.model.device) for k, v in inputs.items()} |
| with torch.no_grad(): |
| out = self.model.generate( |
| **inputs, |
| max_new_tokens=max_new, |
| do_sample=do_sample, |
| temperature=max(temperature, 1e-5) if do_sample else None, |
| pad_token_id=self.tokenizer.pad_token_id, |
| ) |
| text = self.tokenizer.decode( |
| out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True |
| ) |
| text = text.split("<end_of_turn>")[0].strip() |
| return [{"generated_text": text}] |
|
|
