Files changed (1) hide show
  1. handler.py +15 -0
handler.py ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+ import json
3
+
4
+ # Load the model
5
+ def model_fn(model_dir):
6
+ classifier = pipeline("text-classification", model=model_dir)
7
+ return classifier
8
+
9
+ # Handle requests
10
+ def predict_fn(input_data, model):
11
+ # Parse input data
12
+ data = json.loads(input_data)
13
+ texts = data.get("inputs", [])
14
+ predictions = model(texts)
15
+ return predictions