Update handler.py

handler.py

@@ -1,25 +1,34 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-class ModelHandler:
-    def __init__(self):
-        self.tokenizer = None
-        self.model = None
-
-    def initialize(self, model_dir):
-        # Load the tokenizer and model
+class EndpointHandler:
+    def __init__(self, model_dir):
+        # Load tokenizer and model during initialization
         self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
         self.model = AutoModelForCausalLM.from_pretrained(model_dir)
+    
+    def __call__(self, data):
+        """
+        This method processes input data and generates output.
+        :param data: Input data, usually a dictionary with 'inputs' key.
+        """
+        # Extract input prompt
+        inputs = data.get("inputs", "")
+        if not inputs:
+            return {"error": "No input provided"}
 
-    def preprocess(self, inputs):
-        # Preprocess the input prompt
-        return self.tokenizer(inputs, return_tensors="pt", padding=True)
-
-    def inference(self, inputs):
-        # Generate text from the model
-        input_ids = inputs["input_ids"]
-        outputs = self.model.generate(input_ids, max_length=200, temperature=0.7)
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    def postprocess(self, outputs):
-        return {"generated_text": outputs}
+        # Preprocess input
+        encoded_inputs = self.tokenizer(inputs, return_tensors="pt", padding=True)
+        
+        # Generate output
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **encoded_inputs,
+                max_length=200,
+                temperature=0.7,
+                do_sample=True
+            )
+        
+        # Decode and return response
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return {"generated_text": response}