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)