File size: 2,044 Bytes
14280bb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 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) |