Create handler.py

handler.py

@@ -0,0 +1,34 @@
+from typing import Dict, Any
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # Load the tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path)
+        self.model.eval()
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Args:
+            data: A dictionary with the key 'inputs' containing the input text.
+        Returns:
+            A dictionary with the generated text under the key 'generated_text'.
+        """
+        # Extract input text
+        input_text = data.get("inputs", "")
+        if not input_text:
+            return {"error": "No input provided"}
+
+        # Tokenize the input
+        inputs = self.tokenizer(input_text, return_tensors="pt")
+
+        # Generate text
+        with torch.no_grad():
+            outputs = self.model.generate(**inputs, max_length=100)
+
+        # Decode the generated tokens
+        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        return {"generated_text": generated_text}