Create handler.py

handler.py

@@ -0,0 +1,30 @@
+import torch
+import bitsandbytes as bnb
+from transformers import AutoTokenizer, AutoModelForCausalLM  # Adjusted import for the correct model type
+import os
+from typing import Dict, List, Any
+
+# Set the environment variable for the Hugging Face token
+
+
+class EndpointHandler:
+    def __init__(self, model_path="thomasabebe/coolmodel"):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float32)  # Adjusted model class
+
+    def __call__(self, data):
+        # Extract inputs from the data
+        inputs = data.get("inputs", "")
+        
+        # Preprocess the inputs
+        encoded_input = self.tokenizer(inputs, return_tensors='pt', padding=True, truncation=True)
+        
+        # Get predictions
+        output = self.model(**encoded_input)
+        
+        # Postprocess the outputs
+        prediction = output.logits.argmax(-1).tolist()  # Convert the tensor to a list
+        
+        # Return the result
+        return {"label": prediction}
+