Rename handler.py to inference.py

handler.py

@@ -1,27 +0,0 @@
-
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
-
-class TextClassifier:
-    def __init__(self, model_path=""):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
-        self.model = AutoModelForSequenceClassification.from_pretrained(model_path, ignore_mismatched_sizes=True)
-        self.model.eval()
-
-    def predict(self, text):
-        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-            probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-        return {
-            "label": "AI" if torch.argmax(probs) == 1 else "HUMAN",
-            "confidence": probs[0].tolist()
-        }
-
-classifier = None
-
-def predict(inputs):
-    global classifier
-    if classifier is None:
-        classifier = TextClassifier(model_path=".")
-    return classifier.predict(inputs)

inference.py

@@ -0,0 +1,14 @@
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            path,
+            ignore_mismatched_sizes=True  # ضروري لتجاوز المشكلة
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.pipeline = pipeline("text-classification", model=self.model, tokenizer=self.tokenizer)
+
+    def __call__(self, data):
+        inputs = data.get("inputs") if isinstance(data, dict) else data
+        return self.pipeline(inputs)