| """Custom Hugging Face Inference Endpoints handler for ProCreations/grug-27b. |
| |
| Deploy: create an Inference Endpoint from this repo with a GPU holding >=60GB |
| VRAM for bf16 (1x A100-80GB / H100 / H200). The endpoint auto-detects this |
| handler. |
| |
| Request formats accepted: |
| {"inputs": "plain user prompt", "parameters": {...}} |
| {"inputs": {"messages": [{"role": "user", "content": "..."}], |
| "tools": [...optional OpenAI-style tool schemas...]}, |
| "parameters": {"max_new_tokens": 1024, "temperature": 0.6, ...}} |
| |
| Response: |
| [{"generated_text": "<full raw completion>", |
| "think": "<grug reasoning>", "answer": "<final answer>"}] |
| """ |
| from typing import Any, Dict, List |
|
|
| import torch |
| from transformers import AutoModelForImageTextToText, AutoTokenizer |
|
|
| DEFAULTS = {"max_new_tokens": 1024, "temperature": 0.6, "top_p": 0.95, "top_k": 20, |
| "do_sample": True} |
| MAX_NEW_TOKENS_CAP = 8192 |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path: str = ""): |
| self.tokenizer = AutoTokenizer.from_pretrained(path) |
| self.model = AutoModelForImageTextToText.from_pretrained( |
| path, dtype=torch.bfloat16, device_map="auto", |
| attn_implementation="sdpa") |
| self.model.eval() |
|
|
| def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]: |
| inputs = data.get("inputs", "") |
| params = dict(DEFAULTS) |
| params.update(data.get("parameters") or {}) |
| params["max_new_tokens"] = min(int(params.get("max_new_tokens", 1024)), |
| MAX_NEW_TOKENS_CAP) |
| if params.get("temperature", 1.0) <= 0: |
| params["do_sample"] = False |
| params.pop("temperature", None) |
| params.pop("top_p", None) |
| params.pop("top_k", None) |
|
|
| tools = None |
| if isinstance(inputs, dict): |
| messages = inputs.get("messages") or [ |
| {"role": "user", "content": str(inputs.get("text", ""))}] |
| tools = inputs.get("tools") or None |
| elif isinstance(inputs, list): |
| messages = inputs |
| else: |
| messages = [{"role": "user", "content": str(inputs)}] |
|
|
| prompt = self.tokenizer.apply_chat_template( |
| messages, tools=tools, add_generation_prompt=True, tokenize=False) |
| enc = self.tokenizer(prompt, return_tensors="pt").to(self.model.device) |
| gen_kwargs = {k: v for k, v in params.items() |
| if k in ("max_new_tokens", "temperature", "top_p", "top_k", |
| "do_sample", "repetition_penalty", "seed")} |
| gen_kwargs.pop("seed", None) |
| with torch.inference_mode(): |
| out = self.model.generate(**enc, **gen_kwargs, |
| pad_token_id=self.tokenizer.pad_token_id) |
| text = self.tokenizer.decode(out[0][enc["input_ids"].shape[1]:], |
| skip_special_tokens=True) |
| if "</think>" in text: |
| think, answer = text.split("</think>", 1) |
| think = think.replace("<think>", "").strip() |
| answer = answer.strip() |
| else: |
| think, answer = "", text.strip() |
| return [{"generated_text": text, "think": think, "answer": answer}] |
|
|
