|
|
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}] |
|
|
|
