update

Author: bdqnghi

codetf/code_utility/apex/apex_code_utility.py

@@ -1,4 +1,6 @@
 from codetf.code_utility.language_specific_utility import LanguageSpecificUtility
+import lizard
+import re
 
 class ApexCodeUtility(LanguageSpecificUtility):
     def __init__(self):
@@ -104,3 +106,17 @@ def extract_attributes(self, tree, text):
     def get_code_attributes(self, code_snippet):
         tree = self.parse(code_snippet)
         return self.extract_attributes(tree, code_snippet)
+
+    def get_function_content_from_line(self, file_content, line_number):
+        file_content_lines = file_content.split("\n")
+        i = lizard.analyze_file.analyze_source_code("AllTests.java", file_content)
+        
+        function_content = None
+        # print(len(i.function_list))
+        for function_dict in i.function_list:
+            function_dict = function_dict.__dict__
+            if line_number >= function_dict["start_line"] and line_number <= function_dict["end_line"]:
+                function_content_lines = file_content_lines[function_dict["start_line"]-1:function_dict["end_line"]]
+                function_content = "\n".join(function_content_lines)
+                break
+        return function_content

codetf/code_utility/base_utility.py

@@ -1,5 +1,7 @@
 from .ast_parser import ASTParser
-import pyparsing
+import lizard
+import re
+
 class BaseUtility():
 
     def __init__(self, language: str):
@@ -22,3 +24,14 @@ def remove_comments(self, code_snippet):
 
         code_snippet = commentFilter.transformString(code_snippet)
         return code_snippet
+
+
+    def get_method_name(self, function_text):
+        first_line = function_text.splitlines()[0]
+        match = re.search(r"(?<=\s)\w+(?=\()", first_line)
+        function_name = None
+        if match:
+            function_name = match.group()
+        
+        return function_name
+

codetf/code_utility/java/java_code_utility.py

@@ -1,4 +1,6 @@
 from codetf.code_utility.language_specific_utility import LanguageSpecificUtility
+import lizard
+import re
 
 class JavaCodeUtility(LanguageSpecificUtility):
     def __init__(self):
@@ -44,4 +46,18 @@ def transform(self, id_nodes, code_text):
     def rename_identifiers(self, code_snippet):
         tree = self.parse(code_snippet)
         identifier_nodes = self.get_identifier_nodes(tree, code_snippet)
-        return self.transform(identifier_nodes, code_snippet)
+        return self.transform(identifier_nodes, code_snippet)
+
+    def get_function_content_from_line(self, file_content, line_number):
+        file_content_lines = file_content.split("\n")
+        i = lizard.analyze_file.analyze_source_code("AllTests.java", file_content)
+        
+        function_content = None
+        # print(len(i.function_list))
+        for function_dict in i.function_list:
+            function_dict = function_dict.__dict__
+            if line_number >= function_dict["start_line"] and line_number <= function_dict["end_line"]:
+                function_content_lines = file_content_lines[function_dict["start_line"]-1:function_dict["end_line"]]
+                function_content = "\n".join(function_content_lines)
+                break
+        return function_content

codetf/performance/model_evaluator.py

@@ -26,7 +26,7 @@ def evaluate_pass_k(self, problems, unit_tests, batch_size=1, max_length=600,
                         top_p=0.95, k=[1,10,100], 
                         num_return_sequences=200, sequences_per_chunk=10, num_workers=1):
         # Load dataset
-        data_loader = Dat aLoader(problems, batch_size=batch_size)
+        data_loader = DataLoader(problems, batch_size=batch_size)
         data_loader = self.accelerator.prepare(data_loader)
 
         # Initialize stopping criteria