File size: 2,114 Bytes
74f6684 f49875f 7f1720f f49875f 1b4093a f49875f 74f6684 f49875f 74f6684 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
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()
top_guess = None
top_score = None
top_success = False
top_10_guesses = None
top_10_scores = None
top_10_success = False
if len(sorted_programs) > 0:
top_guess = sorted_programs[0][1]
top_score = sorted_programs[0][0]
top_success = parse(top_guess.replace('\\', '')).equivalent(parse(true_program.replace('\\', '')))
top_10_guesses = [p for s, p in sorted_programs[:10]]
top_10_scores = [s for s, p in sorted_programs[:10]]
top_10_success = any([parse(p.replace('\\', '')).equivalent(parse(true_program.replace('\\', ''))) for p in top_10_guesses])
return {
"guess": top_guess,
"top_1_success": top_success,
"top_1_score": top_score,
"top_10_guesses": top_10_guesses,
"top_10_scores": top_10_scores,
"top_10_success": top_10_success,
"time": end - start
}
|