Update handler.py

handler.py

@@ -1,3 +1,23 @@
+import joblib
+import os
+import json
+
+class EndpointHandler:
+    def __init__(self, model_dir):
+        self.model_dir = model_dir
+        self.vectorizer = joblib.load(os.path.join(model_dir, 'vectorizer.joblib'))
+        self.model = joblib.load(os.path.join(model_dir, 'logistic_classifier.joblib'))
+
+        # Verify that the tokenizer configuration is correct
+        with open(os.path.join(model_dir, "tokenizer.json"), "r") as file:
+            tokenizer_config = json.load(file)
+        if tokenizer_config['tokenizer'] != 'split':
+            raise ValueError("Tokenizer configuration does not match the expected tokenizer.")
+
+    def predict_rating(self, review):
+        review_tfidf = self.vectorizer.transform([review])
+        predicted_rating = self.model.predict(review_tfidf)[0]
+        return int(predicted_rating)
 def __call__(self, inputs):
     try:
         inputs_dict = json.loads(inputs)
@@ -29,3 +49,4 @@ def __call__(self, inputs):
 
     except Exception as e:
         return json.dumps({"error": str(e)})
+