| import endpoint_preflight |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from peft import PeftConfig, PeftModel |
| from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path: str = ""): |
| repo_dir = Path(path) |
| for filename in ( |
| "adapter_config.json", "tokenizer.json", "tokenizer_config.json" |
| ): |
| if not (repo_dir / filename).is_file(): |
| raise FileNotFoundError(f"Missing {filename} in {repo_dir}") |
|
|
| config = PeftConfig.from_pretrained( |
| str(repo_dir), local_files_only=True |
| ) |
| if not config.base_model_name_or_path: |
| raise ValueError("The adapter configuration does not identify a base model.") |
|
|
| |
| self.tokenizer = PreTrainedTokenizerFast.from_pretrained( |
| str(repo_dir), local_files_only=True |
| ) |
| if not self.tokenizer.chat_template: |
| raise ValueError("Include the saved GPT-OSS chat template in the repository.") |
|
|
| |
| vocab = self.tokenizer.get_vocab() |
| spellings = ( |
| ("<|start|>", "<|channel|>", "<|message|>", "<|return|>"), |
| ("<|im_start|>", "<|meta_sep|>", "<|im_sep|>", "<|fim_suffix|>"), |
| ) |
| for start, channel, message, finish in spellings: |
| if all(token in vocab for token in (start, channel, message, finish)): |
| self.final_marker = self.tokenizer.encode( |
| channel + "final" + message, add_special_tokens=False |
| ) |
| self.turn_end_id = vocab[finish] |
| self.boundary_ids = {vocab[start], vocab[channel], vocab[message]} |
| break |
| else: |
| raise ValueError("The saved tokenizer is missing GPT-OSS boundary tokens.") |
|
|
| for token in ("<|end|>", "<|im_end|>", "<|call|>", "<|ghissue|>"): |
| if token in vocab: |
| self.boundary_ids.add(vocab[token]) |
| self.boundary_ids.add(self.turn_end_id) |
|
|
| |
| if self.tokenizer.pad_token_id is None: |
| if self.tokenizer.eos_token_id is None: |
| raise ValueError("The saved tokenizer has neither padding nor EOS configured.") |
| self.tokenizer.pad_token = self.tokenizer.eos_token |
|
|
| base_model = AutoModelForCausalLM.from_pretrained( |
| config.base_model_name_or_path, |
| revision=getattr(config, "revision", None), |
| torch_dtype="auto", |
| device_map="auto", |
| ) |
| self.model = PeftModel.from_pretrained( |
| base_model, |
| str(repo_dir), |
| config=config, |
| is_trainable=False, |
| local_files_only=True, |
| ) |
| self.model.eval() |
|
|
| def _final_answer(self, prompt_ids: list[int], new_ids: list[int]) -> str: |
| |
| if not new_ids or new_ids[-1] != self.turn_end_id: |
| raise RuntimeError( |
| "Generation did not finish a complete answer. Shorten the request " |
| "or increase max_new_tokens within the endpoint's memory budget." |
| ) |
|
|
| marker = self.final_marker |
| |
| ids = prompt_ids[-len(marker):] + new_ids |
| positions = [ |
| i + len(marker) |
| for i in range(len(ids) - len(marker) + 1) |
| if ids[i:i + len(marker)] == marker |
| ] |
| if not positions: |
| raise RuntimeError("No GPT-OSS final-answer channel was generated.") |
|
|
| answer_ids = ids[positions[-1]:-1] |
| if any(token in self.boundary_ids for token in answer_ids): |
| raise RuntimeError("Unexpected message boundaries inside the final answer.") |
| answer = self.tokenizer.decode( |
| answer_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False |
| ).strip() |
| if not answer: |
| raise RuntimeError("The model generated an empty final answer.") |
| return answer |
|
|
| def __call__(self, data: dict[str, Any]) -> list[dict[str, str]]: |
| if not isinstance(data, dict): |
| raise ValueError("The request body must be a JSON object.") |
| prompt = data.get("inputs") |
| if not isinstance(prompt, str) or not prompt.strip(): |
| raise ValueError("'inputs' must be a non-empty prompt string.") |
|
|
| params = data.get("parameters", {}) |
| if not isinstance(params, dict): |
| raise ValueError("'parameters' must be a JSON object.") |
| max_new_tokens = params.get("max_new_tokens", 512) |
| if type(max_new_tokens) is not int or not 1 <= max_new_tokens <= 1024: |
| raise ValueError("max_new_tokens must be an integer from 1 to 1024.") |
|
|
| |
| messages = [{"role": "user", "content": prompt}] |
| inputs = self.tokenizer.apply_chat_template( |
| messages, |
| tokenize=True, |
| add_generation_prompt=True, |
| return_tensors="pt", |
| return_dict=True, |
| ) |
| prompt_ids = inputs["input_ids"][0].tolist() |
| context_limit = getattr(self.model.config, "max_position_embeddings", None) |
| if context_limit and len(prompt_ids) + max_new_tokens > context_limit: |
| raise ValueError("The input plus output budget exceeds the model context limit.") |
| inputs = inputs.to(self.model.device) |
|
|
| with torch.inference_mode(): |
| outputs = self.model.generate( |
| input_ids=inputs["input_ids"], |
| attention_mask=inputs["attention_mask"], |
| max_new_tokens=max_new_tokens, |
| do_sample=False, |
| num_beams=1, |
| num_return_sequences=1, |
| return_dict_in_generate=False, |
| pad_token_id=self.tokenizer.pad_token_id, |
| eos_token_id=self.turn_end_id, |
| ) |
|
|
| new_ids = outputs[0][len(prompt_ids):].tolist() |
| return [{"generated_text": self._final_answer(prompt_ids, new_ids)}] |
|
|
