Update RepoPipeline.py

RepoPipeline.py

@@ -10,6 +10,11 @@ from tqdm.auto import tqdm
 
 
 def extract_code_and_docs(text: str):
+    """
+    The method for extracting codes and docs in text.
+    :param text: python file.
+    :return: codes and docs set.
+    """
     code_set = set()
     docs_set = set()
     root = ast.parse(text)
@@ -28,7 +33,33 @@ def extract_code_and_docs(text: str):
     return code_set, docs_set
 
 
+def extract_requirements(lines):
+    """
+    The method for extracting requirements.
+    :param lines: requirements.
+    :return: requirement libraries.
+    """
+    requirements_set = set()
+    for line in lines:
+        try:
+            if line != "\n":
+                if " == " in line:
+                    splitLine = line.split(" == ")
+                else:
+                    splitLine = line.split("==")
+                requirements_set.update(splitLine[0])
+        except:
+            pass
+    return requirements_set
+
+
 def get_metadata(repo_name, headers=None):
+    """
+    The method for getting metadata of repository from github_api.
+    :param repo_name: repository name.
+    :param headers: request headers.
+    :return: response json.
+    """
     api_url = f"https://api.github.com/repos/{repo_name}"
     tqdm.write(f"[+] Getting metadata for {repo_name}")
     try:
@@ -41,9 +72,15 @@ def get_metadata(repo_name, headers=None):
 
 
 def extract_information(repos, headers=None):
+    """
+    The method for extracting repositories information.
+    :param repos: repositories.
+    :param headers: request header.
+    :return: a list for representing the information of each repository.
+    """
     extracted_infos = []
     for repo_name in tqdm(repos, disable=len(repos) <= 1):
-        # Get metadata
+        # 1. Extracting metadata.
         metadata = get_metadata(repo_name, headers=headers)
         repo_info = {
             "name": repo_name,
@@ -60,7 +97,7 @@ def extract_information(repos, headers=None):
         if metadata.get("license"):
             repo_info["license"] = metadata["license"]["spdx_id"]
 
-        # Download repo tarball bytes
+        # Download repo tarball bytes ---- Download repository.
         download_url = f"https://api.github.com/repos/{repo_name}/tarball"
         tqdm.write(f"[+] Downloading {repo_name}")
         try:
@@ -70,24 +107,50 @@ def extract_information(repos, headers=None):
             tqdm.write(f"[-] Failed to download {repo_name}: {e}")
             continue
 
-        # Extract python files and parse them
+        # Extract repository files and parse them
         tqdm.write(f"[+] Extracting {repo_name} info")
         with tarfile.open(fileobj=response.raw, mode="r|gz") as tar:
             for member in tar:
-                if (member.name.endswith(".py") and member.isfile()) is False:
-                    continue
-                try:
-                    file_content = tar.extractfile(member).read().decode("utf-8")
-                    code_set, docs_set = extract_code_and_docs(file_content)
-
-                    repo_info["codes"].update(code_set)
-                    repo_info["docs"].update(docs_set)
-                except UnicodeDecodeError as e:
-                    tqdm.write(
-                        f"[-] UnicodeDecodeError in {member.name}, skipping: \n{e}"
-                    )
-                except SyntaxError as e:
-                    tqdm.write(f"[-] SyntaxError in {member.name}, skipping: \n{e}")
+                # 2. Extracting codes and docs.
+                if (member.name.endswith(".py") and member.isfile()) is True:
+                    try:
+                        file_content = tar.extractfile(member).read().decode("utf-8")
+                        # extract_code_and_docs
+                        code_set, docs_set = extract_code_and_docs(file_content)
+                        repo_info["codes"].update(code_set)
+                        repo_info["docs"].update(docs_set)
+                    except UnicodeDecodeError as e:
+                        tqdm.write(
+                            f"[-] UnicodeDecodeError in {member.name}, skipping: \n{e}"
+                        )
+                    except SyntaxError as e:
+                        tqdm.write(f"[-] SyntaxError in {member.name}, skipping: \n{e}")
+                # 3. Extracting readme.
+                elif ((member.name == "README.md" or member.name == "README.rst") and member.isfile()) is True:
+                    try:
+                        file_content = tar.extractfile(member).read().decode("utf-8")
+                        # extract readme
+                        readme_set = set(file_content)
+                        repo_info["readmes"].update(readme_set)
+                    except UnicodeDecodeError as e:
+                        tqdm.write(
+                            f"[-] UnicodeDecodeError in {member.name}, skipping: \n{e}"
+                        )
+                    except SyntaxError as e:
+                        tqdm.write(f"[-] SyntaxError in {member.name}, skipping: \n{e}")
+                # 4. Extracting requirements.
+                elif (member.name == "requirements.txt" and member.isfile()) is True:
+                    try:
+                        lines = tar.extractfile(member).readlines().decode("utf-8")
+                        # extract readme
+                        requirements_set = extract_requirements(lines)
+                        repo_info["requirements"].update(requirements_set)
+                    except UnicodeDecodeError as e:
+                        tqdm.write(
+                            f"[-] UnicodeDecodeError in {member.name}, skipping: \n{e}"
+                        )
+                    except SyntaxError as e:
+                        tqdm.write(f"[-] SyntaxError in {member.name}, skipping: \n{e}")
 
         extracted_infos.append(repo_info)
 
@@ -95,11 +158,20 @@ def extract_information(repos, headers=None):
 
 
 class RepoPipeline(Pipeline):
+    """
+    A custom pipeline for generating series of embeddings of a repository.
+    """
 
     def __init__(self, github_token=None, *args, **kwargs):
+        """
+        The initial method for pipeline.
+        :param github_token: github_token
+        :param args: args
+        :param kwargs: kwargs
+        """
         super().__init__(*args, **kwargs)
 
-        # Github token
+        # Getting github token
         self.github_token = github_token
         if self.github_token:
             print("[+] GitHub token set!")
@@ -111,36 +183,56 @@ class RepoPipeline(Pipeline):
             )
 
     def _sanitize_parameters(self, **pipeline_parameters):
+        """
+        The method for splitting parameters.
+        :param pipeline_parameters: parameters
+        :return: different parameters of different periods.
+        """
+        # The parameters of "preprocess" period.
         preprocess_parameters = {}
         if "github_token" in pipeline_parameters:
             preprocess_parameters["github_token"] = pipeline_parameters["github_token"]
 
+        # The parameters of "forward" period.
         forward_parameters = {}
         if "max_length" in pipeline_parameters:
             forward_parameters["max_length"] = pipeline_parameters["max_length"]
 
+        # The parameters of "postprocess" period.
         postprocess_parameters = {}
         return preprocess_parameters, forward_parameters, postprocess_parameters
 
     def preprocess(self, input_: Any, github_token=None) -> List:
-        # Making input to list format
+        """
+        The method for "preprocess" period.
+        :param input_: the input.
+        :param github_token: github_token.
+        :return: a list about repository information.
+        """
+        # Making input to list format.
         if isinstance(input_, str):
             input_ = [input_]
 
-        # Building token
+        # Building headers.
         headers = {"Accept": "application/vnd.github+json"}
         token = github_token or self.github_token
         if token:
             headers["Authorization"] = f"Bearer {token}"
 
-        # Getting repositories' information: input_ means series of repositories
+        # Getting repositories' information: input_ means series of repositories (can be only one repository).
         extracted_infos = extract_information(input_, headers=headers)
-
         return extracted_infos
 
     def encode(self, text, max_length):
+        """
+        The method for encoding the text to embedding by using UniXcoder.
+        :param text: text.
+        :param max_length: the max length.
+        :return: the embedding of text.
+        """
         assert max_length < 1024
 
+        # Getting the tokenizer.
         tokenizer = self.tokenizer
         tokens = (
                 [tokenizer.cls_token, "<encoder-only>", tokenizer.sep_token]
@@ -149,20 +241,36 @@ class RepoPipeline(Pipeline):
         )
         tokens_id = tokenizer.convert_tokens_to_ids(tokens)
         source_ids = torch.tensor([tokens_id]).to(self.device)
-
         token_embeddings = self.model(source_ids)[0]
+
+        # Getting the text embedding.
         sentence_embeddings = token_embeddings.mean(dim=1)
 
         return sentence_embeddings
 
     def generate_embeddings(self, text_sets, max_length):
+        """
+        The method for generating embeddings of a text set.
+        :param text_sets: text set.
+        :param max_length: max length.
+        :return: the embeddings of text set.
+        """
         assert max_length < 1024
+
+        # Concat the embeddings of each sentence/text in vertical dimension.
         return torch.zeros((1, 768), device=self.device) \
             if text_sets is None or len(text_sets) == 0 \
             else torch.cat([self.encode(text, max_length) for text in text_sets], dim=0)
 
     def _forward(self, extracted_infos: List, max_length=512) -> List:
+        """
+        The method for "forward" period.
+        :param extracted_infos: the information of repositories.
+        :param max_length: max length.
+        :return: the output of this pipeline.
+        """
         model_outputs = []
+        # The number of repository.
         num_repos = len(extracted_infos)
         with tqdm(total=num_repos) as progress_bar:
             # For each repository
@@ -194,18 +302,37 @@ class RepoPipeline(Pipeline):
                 info["requirement_embeddings"] = requirement_embeddings.cpu().numpy()
                 info["mean_requirement_embedding"] = torch.mean(requirement_embeddings, dim=0).cpu().numpy()
 
-                # Requirement embeddings
+                # Readme embeddings
                 tqdm.write(f"[*] Generating readme embeddings for {repo_name}")
                 readme_embeddings = self.generate_embeddings(repo_info["readmes"], max_length)
                 info["readme_embeddings"] = readme_embeddings.cpu().numpy()
                 info["mean_readme_embedding"] = torch.mean(readme_embeddings, dim=0).cpu().numpy()
 
+                # Repo-level mean embedding
+                info["mean_repo_embedding"] = torch.cat([
+                    info["mean_code_embedding"],
+                    info["mean_doc_embedding"],
+                    info["mean_requirement_embedding"],
+                    info["mean_readme_embedding"]
+                ], dim=1)
+
+                # TODO Remove test
+                info["code_embeddings_shape"] = info["code_embeddings"].shape
+                info["doc_embeddings_shape"] = info["doc_embeddings"].shape
+                info["requirement_embeddings_shape"] = info["requirement_embeddings"].shape
+                info["readme_embeddings_shape"] = info["readme_embeddings"].shape
+                info["mean_repo_embedding_shape"] = info["mean_repo_embedding"].shape
+
                 progress_bar.update(1)
                 model_outputs.append(info)
 
         return model_outputs
 
     def postprocess(self, model_outputs: List, **postprocess_parameters: Dict) -> List:
+        """
+        The method for "postprocess" period.
+        :param model_outputs: the output of this pipeline.
+        :param postprocess_parameters: the parameters of "postprocess" period.
+        :return: model output.
+        """
         return model_outputs
-
-