add handler.py

handler.py

@@ -0,0 +1,16 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from typing import Dict
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path)
+
+    def __call__(self, data: Dict[str, str]) -> Dict[str, str]:
+        inputs = data.get("inputs", "")
+        if not inputs:
+            return {"error": "No input provided."}
+        inputs = self.tokenizer(inputs, return_tensors="pt")
+        outputs = self.model.generate(**inputs, max_new_tokens=100)
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return {"generated_text": response}

load_model.py

@@ -1,6 +0,0 @@
-from transformers import AutoModel, AutoTokenizer
-
-def load_custom_model():
-    tokenizer = AutoTokenizer.from_pretrained("./model")
-    model = AutoModel.from_pretrained("./model")
-    return tokenizer, model