Update agent.py

princeton-vl · Oct 10, 2020 · 4c205a7 · 4c205a7
1 parent 85d3c85
commit 4c205a7
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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