From 4c205a78485b4c7dd355441d816d299bc984998a Mon Sep 17 00:00:00 2001 From: Kaiyu Yang Date: Sat, 10 Oct 2020 10:37:11 -0400 Subject: [PATCH] Update agent.py --- ASTactic/agent.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/ASTactic/agent.py b/ASTactic/agent.py index 37d3862d..9fd3c516 100644 --- a/ASTactic/agent.py +++ b/ASTactic/agent.py @@ -32,8 +32,10 @@ def action_seq_loss(logits_batch, actions_batch, opts): sexp_cache = SexpCache('../sexp_cache', readonly=True) def filter_env(env): + "Get the last 10 toplevel constants" filtered_env = [] - for const in [const for const in env['constants'] if const['qualid'].startswith('SerTop')][-10:]: + toplevel_consts = [const for const in env['constants'] if const['qualid'].startswith('SerTop')] + for const in toplevel_consts[-10:]: ast = sexp_cache[const['sexp']] filtered_env.append({'qualid': const['qualid'], 'ast': term_parser.parse(ast)}) return filtered_env