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