Neural Coder fix bugs (#394)

* Update handle_user_input.py

* Update function.py

* Update code_line.py

Author: kaikaiyao

neural_coder/graphers/code_line.py

@@ -152,8 +152,11 @@ def register_code_line():
                     # judge_1: indent is equal to def indent
                     judge_1 = following_indent_level <= class_def_indent_level
                     # judge_2: not starting with")"
-                    judge_2 = True if (
-                        following_line != "" and following_line[following_indent_level] != ")") else False
+                    try:
+                        judge_2 = True if (
+                            following_line != "" and following_line[following_indent_level] != ")") else False
+                    except:
+                        judge_2 = False
                     # judge_3: is not a comment or empty line
                     judge_3 = True if (
                         not _is_multi_line_comment and not _is_single_line_comment_or_empty) else False
@@ -216,8 +219,11 @@ def register_code_line():
                     # judge_1: indent is equal to def indent
                     judge_1 = following_indent_level <= func_def_indent_level
                     # judge_2: not starting with")"
-                    judge_2 = True if (
-                        following_line != "" and following_line[following_indent_level] != ")") else False
+                    try:
+                        judge_2 = True if (
+                            following_line != "" and following_line[following_indent_level] != ")") else False
+                    except:
+                        judge_2 = False
                     # judge_3: is not a comment or empty line
                     judge_3 = True if (
                         not _is_multi_line_comment and not _is_single_line_comment_or_empty) else False

neural_coder/graphers/function.py

@@ -79,8 +79,11 @@ def c():
                     # judge_1: indent is equal to def indent
                     judge_1 = following_indent_level <= def_indent_level
                     # judge_2: not starting with")"
-                    judge_2 = True if (
-                        following_line != "" and following_line[following_indent_level] != ")") else False
+                    try:
+                        judge_2 = True if (
+                            following_line != "" and following_line[following_indent_level] != ")") else False
+                    except:
+                        judge_2 = False
                     # judge_3: is not a comment
                     c1 = False
                     c2 = False

neural_coder/utils/handle_user_input.py

@@ -100,8 +100,11 @@ def get_imports_path(user_code_path: List) -> List:
         "classy_vision",
         "collections",
         "contextlib",
-        "datasets",
+        "data",
         "dataclasses",
+        "dataloader",
+        "dataset",
+        "datasets",
         "einops",
         "enum",
         "evaluate",
@@ -112,13 +115,17 @@ def get_imports_path(user_code_path: List) -> List:
         "importlib",
         "inspect",
         "intel_extension_for_pytorch",
+        "itertools",
         "jax",
         "json",
         "keras",
         "libcst",
         "linecache",
         "logging",
+        "loggings",
         "math",
+        "model",
+        "models",
         "neural_coder",
         "neural_compressor",
         "numpy",
@@ -131,6 +138,7 @@ def get_imports_path(user_code_path: List) -> List:
         "shutil",
         "subprocess",
         "sys",
+        "system",
         "tempfile",
         "typing",
         "tensorflow",
@@ -141,11 +149,15 @@ def get_imports_path(user_code_path: List) -> List:
         "torchdynamo",
         "tqdm",
         "traceback",
+        "trainer_qa",
         "unittest",
+        "utils_qa",
         "vissl",
+        "warning",
         "warnings",
         "packaging",
         "pytorch_lightning",
+        "__future__",
     ]
 
     if globals.cache_load_transformers: