Create handler.py

14280bb verified about 2 years ago

2.04 kB

import os
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import torch



class HuggingFaceHandler:
def init(self, model_dir: str, task: str):
"""
Initialize the HuggingFaceHandler with the model directory and task.
This loads the model and tokenizer.
"""
self.model_dir = model_dir
self.task = task



# Load the model and tokenizer without specifying a device
self.pipeline = self.get_pipeline()



def get_pipeline(self):
"""
Loads the model and tokenizer and sets up the Hugging Face pipeline.
Let accelerate handle device placement, and remove any device argument.
"""
try:
# Load the tokenizer and model from the specified directory
tokenizer = AutoTokenizer.from_pretrained(self.model_dir)
model = AutoModelForSeq2SeqLM.from_pretrained(self.model_dir)



# Create the Hugging Face pipeline (without specifying device)
hf_pipeline = pipeline(task=self.task, model=model, tokenizer=tokenizer)
return hf_pipeline
except Exception as e:
raise RuntimeError(f"Error loading model and tokenizer: {str(e)}")



def predict(self, inputs: str) -> str:
"""
Make predictions using the pipeline.
:param inputs: Text input for prediction
:return: Generated text or task-specific output
"""
try:
# Pass the input text to the pipeline and generate the output
result = self.pipeline(inputs)
return result
except Exception as e:
raise RuntimeError(f"Error during inference: {str(e)}")



def get_inference_handler_either_custom_or_default_handler(model_dir: str, task: str):
"""
Helper function to return the handler instance.
"""
return HuggingFaceHandler(model_dir=model_dir, task=task)



# Example of usage
if name == "__main__":
model_directory = os.getenv("MODEL_DIR", "maxmnd/fine_tuned")
task_type = os.getenv("TASK", "text-generation")



# Instantiate the handler
handler = get_inference_handler_either_custom_or_default_handler(model_directory, task_type)



# Example input text
input_text = "What is the capital of France?"



# Perform inference
output = handler.predict(input_text)
print("Generated Output:", output)

	import os
	from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
	import torch



	class HuggingFaceHandler:
	def init(self, model_dir: str, task: str):
	"""
	Initialize the HuggingFaceHandler with the model directory and task.
	This loads the model and tokenizer.
	"""
	self.model_dir = model_dir
	self.task = task



	# Load the model and tokenizer without specifying a device
	self.pipeline = self.get_pipeline()



	def get_pipeline(self):
	"""
	Loads the model and tokenizer and sets up the Hugging Face pipeline.
	Let accelerate handle device placement, and remove any device argument.
	"""
	try:
	# Load the tokenizer and model from the specified directory
	tokenizer = AutoTokenizer.from_pretrained(self.model_dir)
	model = AutoModelForSeq2SeqLM.from_pretrained(self.model_dir)



	# Create the Hugging Face pipeline (without specifying device)
	hf_pipeline = pipeline(task=self.task, model=model, tokenizer=tokenizer)
	return hf_pipeline
	except Exception as e:
	raise RuntimeError(f"Error loading model and tokenizer: {str(e)}")



	def predict(self, inputs: str) -> str:
	"""
	Make predictions using the pipeline.
	:param inputs: Text input for prediction
	:return: Generated text or task-specific output
	"""
	try:
	# Pass the input text to the pipeline and generate the output
	result = self.pipeline(inputs)
	return result
	except Exception as e:
	raise RuntimeError(f"Error during inference: {str(e)}")



	def get_inference_handler_either_custom_or_default_handler(model_dir: str, task: str):
	"""
	Helper function to return the handler instance.
	"""
	return HuggingFaceHandler(model_dir=model_dir, task=task)



	# Example of usage
	if name == "__main__":
	model_directory = os.getenv("MODEL_DIR", "maxmnd/fine_tuned")
	task_type = os.getenv("TASK", "text-generation")



	# Instantiate the handler
	handler = get_inference_handler_either_custom_or_default_handler(model_directory, task_type)



	# Example input text
	input_text = "What is the capital of France?"



	# Perform inference
	output = handler.predict(input_text)
	print("Generated Output:", output)