Create handler.py

handler.py

@@ -0,0 +1,85 @@
+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)