Create handler.py

handler.py

@@ -0,0 +1,79 @@
+import json
+from bertopic import BERTopic
+
+class EndpointHandler:
+    def __init__(self, model_path="SCANSKY/BERTopic-Tourism-English"):
+        """
+        Initialize the handler. Load the BERTopic model from Hugging Face.
+        """
+        self.topic_model = BERTopic.load(model_path)
+
+    def preprocess(self, data):
+        """
+        Preprocess the incoming request data.
+        - Extract text input from the request.
+        """
+        try:
+            # Parse the incoming data
+            body = json.loads(data)
+            text_input = body.get("text", "")
+            return text_input
+        except Exception as e:
+            raise ValueError(f"Error during preprocessing: {str(e)}")
+
+    def inference(self, text_input):
+        """
+        Perform inference using the BERTopic model.
+        - Process the text input and generate topic predictions.
+        """
+        try:
+            # Split text into documents (assuming one document per line)
+            docs = text_input.strip().split('\n')
+
+            # Perform topic inference
+            topics, probabilities = self.topic_model.transform(docs)
+
+            # Get topic information
+            topics_info = self.topic_model.get_topics()
+
+            # Prepare the results
+            results = []
+            for topic, prob in zip(topics, probabilities):
+                topic_words = topics_info.get(topic, [])
+                results.append({
+                    "topic": int(topic),
+                    "probability": float(prob),
+                    "top_words": [word for word, _ in topic_words[:5]]  # Top 5 words
+                })
+
+            return results
+        except Exception as e:
+            raise ValueError(f"Error during inference: {str(e)}")
+
+    def postprocess(self, results):
+        """
+        Postprocess the inference results into a JSON response.
+        """
+        return {
+            "results": results
+        }
+
+    def __call__(self, data):
+        """
+        Handle the incoming request.
+        """
+        try:
+            # Preprocess the data
+            text_input = self.preprocess(data)
+
+            # Perform inference
+            results = self.inference(text_input)
+
+            # Postprocess the results
+            response = self.postprocess(results)
+
+            return response
+        except Exception as e:
+            return {
+                "error": str(e)
+            }