|
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): |
|
|
|
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 |
|
} |
|
|
