timgpt / handler.py

Update handler.py

65c065f verified almost 2 years ago

1.59 kB

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch

class EndpointHandler:
    def __init__(self, path="krisoei/timgpt"):
        if not path:
            raise ValueError("A valid model path or name must be provided.")
        
        # Load tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForCausalLM.from_pretrained(
            path, 
            torch_dtype=torch.float16, 
            device_map="auto"
        )
        
        # Set up text-generation pipeline
        self.pipe = pipeline(
            "text-generation",
            model=self.model,
            tokenizer=self.tokenizer,
            max_new_tokens=512,
            do_sample=True,
            temperature=0.7,
            top_p=0.95,
        )

    def __call__(self, data):
        # Validate input data
        if not isinstance(data, dict):
            return {"error": "Input must be a JSON object."}
        
        prompt = data.get("inputs", "")
        if not prompt:
            return {"error": "No input provided."}
        
        try:
            # Generate response
            outputs = self.pipe(prompt)
            if outputs:
                response = outputs[0]['generated_text']
                # Remove the original prompt from the response
                response = response[len(prompt):].strip()
                return {"generated_text": response}
            else:
                return {"error": "No output generated."}
        except Exception as e:
            return {"error": str(e)}

	from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
	import torch

	class EndpointHandler:
	def __init__(self, path="krisoei/timgpt"):
	if not path:
	raise ValueError("A valid model path or name must be provided.")

	# Load tokenizer and model
	self.tokenizer = AutoTokenizer.from_pretrained(path)
	self.model = AutoModelForCausalLM.from_pretrained(
	path,
	torch_dtype=torch.float16,
	device_map="auto"
	)

	# Set up text-generation pipeline
	self.pipe = pipeline(
	"text-generation",
	model=self.model,
	tokenizer=self.tokenizer,
	max_new_tokens=512,
	do_sample=True,
	temperature=0.7,
	top_p=0.95,
	)

	def __call__(self, data):
	# Validate input data
	if not isinstance(data, dict):
	return {"error": "Input must be a JSON object."}

	prompt = data.get("inputs", "")
	if not prompt:
	return {"error": "No input provided."}

	try:
	# Generate response
	outputs = self.pipe(prompt)
	if outputs:
	response = outputs[0]['generated_text']
	# Remove the original prompt from the response
	response = response[len(prompt):].strip()
	return {"generated_text": response}
	else:
	return {"error": "No output generated."}
	except Exception as e:
	return {"error": str(e)}