Update handler.py

8a4e9e9 verified 8 months ago

2.45 kB

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

MODEL_NAME = "teapotai/tinyteapot"
MAX_INPUT_TOKENS = 512


class EndpointHandler:
    def __init__(self, path: str = ""):
        # EXACTLY like your snippet: load by repo id from Hugging Face Hub
        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

        self.model.eval()
        self.device = torch.device("cpu")
        self.model.to(self.device)

        self.system_prompt = (
            "You are Teapot, an open-source AI assistant optimized for low-end devices, "
            "providing short, accurate responses without hallucinating while excelling at "
            "information extraction and text summarization. "
            "If the context does not answer the question, reply exactly: "
            "'I am sorry but I don't have any information on that'."
        )

    @torch.inference_mode()
    def __call__(self, data):
        inputs = data.get("inputs")
        if inputs is None:
            raise ValueError("Missing required field 'inputs'.")

        # Support:
        # 1) inputs is a full prompt string (closest to your snippet)
        # 2) inputs is {"context": ..., "question": ...}
        if isinstance(inputs, str):
            prompt = inputs
        elif isinstance(inputs, dict):
            context = inputs.get("context", "")
            question = inputs.get("question", "")
            prompt = f"{context}\n{self.system_prompt}\n{question}\n"
        else:
            raise ValueError("inputs must be a string OR an object with 'context' and 'question'.")

        # Match your code path
        enc = self.tokenizer(prompt, return_tensors="pt")
        input_ids = enc["input_ids"]
        attention_mask = enc["attention_mask"]

        # Keep only the most recent tokens if too long
        if input_ids.shape[1] > MAX_INPUT_TOKENS:
            input_ids = input_ids[:, -MAX_INPUT_TOKENS:]
            attention_mask = attention_mask[:, -MAX_INPUT_TOKENS:]

        input_ids = input_ids.to(self.device)
        attention_mask = attention_mask.to(self.device)

        outputs = self.model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            do_sample=False,
        )

        answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return {"generated_text": answer}

	from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
	import torch

	MODEL_NAME = "teapotai/tinyteapot"
	MAX_INPUT_TOKENS = 512


	class EndpointHandler:
	def __init__(self, path: str = ""):
	# EXACTLY like your snippet: load by repo id from Hugging Face Hub
	self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
	self.model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

	self.model.eval()
	self.device = torch.device("cpu")
	self.model.to(self.device)

	self.system_prompt = (
	"You are Teapot, an open-source AI assistant optimized for low-end devices, "
	"providing short, accurate responses without hallucinating while excelling at "
	"information extraction and text summarization. "
	"If the context does not answer the question, reply exactly: "
	"'I am sorry but I don't have any information on that'."
	)

	@torch.inference_mode()
	def __call__(self, data):
	inputs = data.get("inputs")
	if inputs is None:
	raise ValueError("Missing required field 'inputs'.")

	# Support:
	# 1) inputs is a full prompt string (closest to your snippet)
	# 2) inputs is {"context": ..., "question": ...}
	if isinstance(inputs, str):
	prompt = inputs
	elif isinstance(inputs, dict):
	context = inputs.get("context", "")
	question = inputs.get("question", "")
	prompt = f"{context}\n{self.system_prompt}\n{question}\n"
	else:
	raise ValueError("inputs must be a string OR an object with 'context' and 'question'.")

	# Match your code path
	enc = self.tokenizer(prompt, return_tensors="pt")
	input_ids = enc["input_ids"]
	attention_mask = enc["attention_mask"]

	# Keep only the most recent tokens if too long
	if input_ids.shape[1] > MAX_INPUT_TOKENS:
	input_ids = input_ids[:, -MAX_INPUT_TOKENS:]
	attention_mask = attention_mask[:, -MAX_INPUT_TOKENS:]

	input_ids = input_ids.to(self.device)
	attention_mask = attention_mask.to(self.device)

	outputs = self.model.generate(
	input_ids=input_ids,
	attention_mask=attention_mask,
	do_sample=False,
	)

	answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
	return {"generated_text": answer}