Use ast to generate the hash of Paddle models.

Author: Xreki

graph_net/paddle/validate.py

@@ -11,7 +11,7 @@
 import numpy as np
 import graph_net
 import os
-import re
+import ast
 import paddle
 
 
@@ -31,26 +31,32 @@ def _get_sha_hash(content):
 
 def _save_to_model_path(dump_dir, hash_text):
     file_path = f"{dump_dir}/graph_hash.txt"
+    print(f"Writing to {file_path}.")
     with open(file_path, "w") as f:
         f.write(hash_text)
 
 
-def extract_from_forward_regex(text, case_sensitive=True):
-    pattern = r"forward.*"
-    flags = 0 if case_sensitive else re.IGNORECASE
+def extract_forward_source(model_path):
+    source = None
+    with open(f"{model_path}/model.py", "r") as f:
+        source = f.read()
 
-    match = re.search(pattern, text, flags)
-    if match:
-        return match.group(0)
-    else:
-        raise ValueError("Erroneous case occurs.")
+    tree = ast.parse(source)
+    forward_code = None
+
+    for node in tree.body:
+        if isinstance(node, ast.ClassDef) and node.name == "GraphModule":
+            for fn in node.body:
+                if isinstance(fn, ast.FunctionDef) and fn.name == "forward":
+                    return ast.unparse(fn)
+    return None
 
 
 def main(args):
     model_path = args.model_path
-    with open(f"{model_path}/model.py", "r") as fp:
-        model_str = fp.read()
-        model_str = extract_from_forward_regex(model_str)
+    if args.dump_graph_hash_key:
+        model_str = extract_forward_source(model_path)
+        assert model_str is not None, f"model_str of {args.model_path} is None."
         _save_to_model_path(model_path, _get_sha_hash(model_str))
 
     model_path = args.model_path
@@ -100,17 +106,16 @@ def main(args):
         required=True,
         help="Path to folder e.g '../test_dataset'",
     )
-
     parser.add_argument(
         "--no-check-redundancy",
         action="store_true",
+        default=False,
         help="whether check model graph redundancy",
     )
-
     parser.add_argument(
         "--dump-graph-hash-key",
         action="store_true",
-        default=False,
+        default=True,
         help="Dump graph hash key",
     )
     parser.add_argument(