oleksandrfluxon
/

mpt-7b-instruct-evaluate

Text Generation

text-generation-inference

Model card Files Files and versions Community

mpt-7b-instruct-evaluate / _handler.py

oleksandrfluxon's picture

oleksandrfluxon

Rename handler.py to _handler.py

0b6b920 11 months ago

raw history blame contribute delete

No virus

2.68 kB

	import torch

	from typing import Any, Dict
	from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
	from accelerate import dispatch_model, infer_auto_device_map
	from accelerate.utils import get_balanced_memory


	class EndpointHandler:
	def __init__(self, path=""):
	with torch.autocast('cuda'):
	path = "oleksandrfluxon/mpt-7b-instruct-evaluate"
	config = AutoConfig.from_pretrained(
	path,
	trust_remote_code=True
	)
	# config.attn_config['attn_impl'] = 'triton'
	config.init_device = 'cuda:0' # For fast initialization directly on GPU!
	config.max_seq_len = 4096 # (input + output) tokens can now be up to 4096


	# load model and tokenizer from path
	self.tokenizer = AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b', padding_side="left")
	model = AutoModelForCausalLM.from_pretrained(
	path,
	config,
	device_map="auto",
	torch_dtype=torch.float16,
	trust_remote_code=True,
	load_in_8bit=True # Load model in the lowest 4-bit precision quantization
	)

	max_memory = get_balanced_memory(
	model,
	max_memory=None,
	no_split_module_classes=["MPTBlock"],
	dtype='float16',
	low_zero=False
	)
	device_map = infer_auto_device_map(
	model,
	max_memory=max_memory,
	no_split_module_classes=["MPTBlock"],
	dtype='float16'
	)
	self.model = dispatch_model(model, device_map=device_map)

	self.device = "cuda" if torch.cuda.is_available() else "cpu"

	def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
	# process input
	inputs = data.pop("inputs", data)
	parameters = data.pop("parameters", None)

	with torch.autocast('cuda'):
	# preprocess
	# inputs = self.tokenizer(inputs, return_tensors="pt").to(self.device)
	inputs = self.tokenizer(inputs, return_tensors="pt")

	# pass inputs with all kwargs in data
	if parameters is not None:
	outputs = self.model.generate(inputs, parameters)
	else:
	outputs = self.model.generate(**inputs)

	# postprocess the prediction
	prediction = self.tokenizer.decode(outputs[0], skip_special_tokens=True)

	return [{"generated_text": prediction}]