File size: 1,337 Bytes
37bde29 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
from typing import Any
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
def extract_assistant_response_simple(response_text):
# Split by assistant header and eot
parts = response_text.split("<|start_header_id|>assistant<|end_header_id|>")[
1
].split("<|eot_id|>")[0]
return parts.strip()
class EndpointHandler:
def __init__(self, path=""):
self.tokenizer = AutoTokenizer.from_pretrained(path)
model = AutoModelForCausalLM.from_pretrained(
path,
torch_dtype="auto",
)
self.model = model.to_bettertransformer()
def __call__(self, data: Any):
text = data.pop("inputs", data)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": text},
]
inputs = self.tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True, # Must add for generation
return_tensors="pt",
).to("cuda")
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
outputs = self.model.generate(**inputs)
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|