984ca58057
Model: serving-d-cause/writing-roleplay-20k-context-nemo-12b-v1.0 Source: Original Platform
43 lines
1.8 KiB
Python
43 lines
1.8 KiB
Python
import os
|
|
from typing import Any, Dict, List
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
|
|
|
|
# Ensure your template includes escaped newlines correctly
|
|
CHAT_TEMPLATE = (
|
|
"{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}"
|
|
"{% for message in messages %}"
|
|
"{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n'}}"
|
|
"{% endfor %}"
|
|
"{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}"
|
|
)
|
|
|
|
class EndpointHandler:
|
|
def __init__(self, model_dir: str, **kwargs: Any):
|
|
# Load tokenizer and model from provided model_dir
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
|
|
# Assign the chat_template under 'default'
|
|
self.tokenizer.chat_template = {"default": CHAT_TEMPLATE}
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|
model_dir, trust_remote_code=True, device_map="auto"
|
|
)
|
|
self.pipeline = TextGenerationPipeline(
|
|
model=self.model, tokenizer=self.tokenizer, return_full_text=False
|
|
)
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
|
# Validate input structure
|
|
messages: List[Dict[str, str]] = data.get("messages")
|
|
if messages is None:
|
|
raise ValueError("Request body must include 'messages' array.")
|
|
|
|
# Format prompt with proper controlling flag
|
|
inputs = self.tokenizer.apply_chat_template(
|
|
messages=messages, add_generation_prompt=True, return_tensors=None
|
|
)
|
|
# Generate text
|
|
gen = self.pipeline(inputs, max_new_tokens=data.get("parameters", {}).get("max_new_tokens", 128))
|
|
content = gen[0]["generated_text"]
|
|
|
|
return {"choices": [{"message": {"role": "assistant", "content": content}}]}
|
|
|