| import time | |
| import torch | |
| import random | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| class EndpointHandler: | |
| """ | |
| Custom handler for `Qwen/Qwen2.5-Math-7B-Instruct`. | |
| """ | |
| def __init__(self, path=""): | |
| """ | |
| Initialize model and tokenizer. | |
| :param path: Path to model and tokenizer | |
| """ | |
| self.tokenizer = AutoTokenizer.from_pretrained(path) | |
| self.model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.bfloat16, device_map="auto") | |
| def __call__(self, data: dict): | |
| """ | |
| Execute model based on input data. | |
| :param data: Input parameters for the model. | |
| Should be in the following form: | |
| `{"inputs": "input_string", "parameters": {"parameter_1": 0, "parameter_2": 0}}` | |
| :return: dict (answer, num_new_token, speed) | |
| """ | |
| question = data.get("inputs", None) | |
| max_new_tokens = data.get("max_new_tokens", 1024) | |
| parameters = data.get("parameters", {}) | |
| if not question: | |
| raise ValueError("Input prompt is missing.") | |
| messages = [ | |
| {"role": "system", "content": "Please reason step by step, and put your final answer within \\boxed{}. " | |
| "Then, give your confidence level in percentage regarding your answer."}, | |
| {"role": "user", "content": question} | |
| ] | |
| tokenized_prompt = self.tokenizer.apply_chat_template( | |
| messages, | |
| add_generation_prompt=True, | |
| return_tensors="pt" | |
| ).to("cuda") | |
| torch.manual_seed(random.randint(0, 2 ** 32 - 1)) | |
| time_start = time.time() | |
| out = self.model.generate( | |
| tokenized_prompt, | |
| max_new_tokens=max_new_tokens, | |
| temperature=1.0, | |
| do_sample=True, | |
| top_p=0.9, | |
| **parameters | |
| ) | |
| time_end = time.time() | |
| response = self.tokenizer.decode(out[0][len(tokenized_prompt[0]):]) | |
| num_new_tokens = len(out[0]) - len(tokenized_prompt[0]) | |
| speed = num_new_tokens / (time_end - time_start) | |
| return { | |
| "answer": response, | |
| "num_new_tokens": num_new_tokens, | |
| "speed": speed | |
| } | |