handler.py

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)