Create handler.py

handler.py

@@ -0,0 +1,38 @@
+from transformers import pipeline
+from greenery import parse
+from greenery.parse import NoMatch
+from listener import Listener, ListenerOutput
+import time
+import json
+import torch
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.listener = Listener(path, {
+            "do_sample": True,
+            "max_new_tokens": 128,
+            "top_p": 0.9,
+            "num_return_sequences": 500,
+            "num_beams": 1
+        }, device="cuda" if torch.cuda.is_available() else "cpu")
+
+    def __call__(self, data):
+        # get inputs
+        inp = data.pop("inputs", None)
+        spec = inp["spec"]
+        true_program = inp["true_program"]
+        
+        start = time.time()
+        outputs = self.listener.synthesize([[(s["string"], s["label"]) for s in spec]], return_scores=True)
+        consistent_program_scores = [outputs.decoded_scores[0][i] for i in outputs.idx[0]]
+        consistent_programs = [outputs.decoded[0][i] for i in outputs.idx[0]]
+        sorted_programs = sorted(set(zip(consistent_program_scores, consistent_programs)), reverse=True, key=lambda x: x[0])
+        end = time.time()
+
+        return {
+            "guess": sorted_programs[0][1],
+            "top_1_success": parse(sorted_programs[0][1]).equivalent(parse(true_program)),
+            "top_1_score": sorted_programs[0][0],
+            "top_5_success": any([parse(p).equivalent(parse(true_program)) for _, p in sorted_programs[:5]]),
+            "time": end - start
+        }