Upload handler.py

handler.py

@@ -1,28 +1,29 @@
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Load model and tokenizer
-model_name = "Syko/SykoNaught-v1"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+class EndpointHandler:
+    def __init__(self, path):
+        # Load tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path)
 
-def handle(inputs):
-    """
-    Handle incoming inference requests.
-    """
-    input_text = inputs.get("inputs", "")
-    max_new_tokens = inputs.get("parameters", {}).get("max_new_tokens", 50)
-    temperature = inputs.get("parameters", {}).get("temperature", 0.7)
+    def __call__(self, inputs):
+        # Parse input
+        input_text = inputs.get("inputs", "")
+        parameters = inputs.get("parameters", {})
+        max_new_tokens = parameters.get("max_new_tokens", 50)
+        temperature = parameters.get("temperature", 0.7)
 
-    # Tokenize input
-    input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+        # Tokenize input
+        input_ids = self.tokenizer(input_text, return_tensors="pt").input_ids
 
-    # Generate output
-    output = model.generate(
-        input_ids,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-    )
-    output_text = tokenizer.decode(output[0], skip_special_tokens=True)
+        # Generate output
+        output = self.model.generate(
+            input_ids,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+        )
 
-    return {"generated_text": output_text}
+        # Decode output
+        output_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
+        return {"generated_text": output_text}