| | from typing import Dict, List, Any |
| | from transformers import AutoModelForCausalLM, AutoTokenizer |
| | import torch |
| |
|
| | class EndpointHandler(): |
| | def __init__(self, path=""): |
| | """ |
| | Initialize the model and tokenizer using the local path. |
| | Uses Zenith Coder v1.1 custom code (modeling_deepseek.py, configuration_deepseek.py, tokenization_deepseek_fast.py). |
| | """ |
| | self.tokenizer = AutoTokenizer.from_pretrained( |
| | path, trust_remote_code=True |
| | ) |
| | self.model = AutoModelForCausalLM.from_pretrained( |
| | path, |
| | trust_remote_code=True, |
| | torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, |
| | device_map="auto" |
| | ) |
| | self.model.eval() |
| |
|
| | def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
| | """ |
| | Accepts a dictionary with the prompt and optional `max_new_tokens`. |
| | Returns generated text. |
| | """ |
| | prompt = data.get("inputs") or data.get("prompt") |
| | if not prompt or not isinstance(prompt, str): |
| | return [{"error": "No valid input prompt provided."}] |
| | |
| | max_new_tokens = int(data.get("max_new_tokens", 256)) |
| | temperature = float(data.get("temperature", 1.0)) |
| | top_p = float(data.get("top_p", 0.95)) |
| | top_k = int(data.get("top_k", 50)) |
| |
|
| | input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids |
| | if torch.cuda.is_available(): |
| | input_ids = input_ids.cuda() |
| | |
| | with torch.no_grad(): |
| | generated_ids = self.model.generate( |
| | input_ids, |
| | do_sample=True, |
| | max_new_tokens=max_new_tokens, |
| | 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 |
| | ) |
| | |
| | gen_text = self.tokenizer.decode( |
| | generated_ids[0][input_ids.shape[1]:], |
| | skip_special_tokens=True |
| | ) |
| | return [{"generated_text": gen_text}] |
| |
|
