Update RepoPipeline.py

RepoPipeline.py

@@ -2,20 +2,14 @@ from typing import Dict, Any, List
 
 import ast
 import tarfile
+from ast import AsyncFunctionDef, ClassDef, FunctionDef, Module
 import torch
 import requests
-import numpy as np
-from ast import AsyncFunctionDef, ClassDef, FunctionDef, Module
 from transformers import Pipeline
 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)
@@ -34,33 +28,7 @@ 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.add(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:
@@ -73,15 +41,9 @@ 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):
-        # 1. Extracting metadata.
+        # Get metadata
         metadata = get_metadata(repo_name, headers=headers)
         repo_info = {
             "name": repo_name,
@@ -98,7 +60,7 @@ def extract_information(repos, headers=None):
         if metadata.get("license"):
             repo_info["license"] = metadata["license"]["spdx_id"]
 
-        # Download repo tarball bytes ---- Download repository.
+        # Download repo tarball bytes
         download_url = f"https://api.github.com/repos/{repo_name}/tarball"
         tqdm.write(f"[+] Downloading {repo_name}")
         try:
@@ -108,51 +70,24 @@ def extract_information(repos, headers=None):
             tqdm.write(f"[-] Failed to download {repo_name}: {e}")
             continue
 
-        # Extract repository files and parse them
+        # Extract python 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:
-                # 2. Extracting codes and docs.
-                if member.name.endswith(".py") and member.isfile():
-                    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():
-                    try:
-                        file_content = tar.extractfile(member).read().decode("utf-8")
-                        # extract readme
-                        readmes_set = set()
-                        readmes_set.add(file_content)
-                        repo_info["readmes"].update(readmes_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():
-                    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}")
+                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}")
 
         extracted_infos.append(repo_info)
 
@@ -160,20 +95,11 @@ 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)
 
-        # Getting github token
+        # Github token
         self.github_token = github_token
         if self.github_token:
             print("[+] GitHub token set!")
@@ -185,56 +111,36 @@ 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:
-        """
-        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.
+        # Making input to list format
         if isinstance(input_, str):
             input_ = [input_]
 
-        # Building headers.
+        # Building token
         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 (can be only one repository).
+        # Getting repositories' information: input_ means series of repositories
         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]
@@ -243,36 +149,20 @@ 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.
+        token_embeddings = self.model(source_ids)[0]
         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 not text_sets \
+            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
@@ -304,26 +194,14 @@ class RepoPipeline(Pipeline):
                 info["requirement_embeddings"] = requirement_embeddings.cpu().numpy()
                 info["mean_requirement_embedding"] = torch.mean(requirement_embeddings, dim=0).cpu().numpy()
 
-                # Readme embeddings
+                # Requirement 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"] = np.concatenate([
-                    info["mean_code_embedding"],
-                    info["mean_doc_embedding"],
-                    info["mean_requirement_embedding"],
-                    info["mean_readme_embedding"]
-                ], axis=0)
-
-                # 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
+                info["doc_embeddings_shape"] = info["doc_embeddings_shape"].shape
 
                 progress_bar.update(1)
                 model_outputs.append(info)
@@ -331,10 +209,6 @@ class RepoPipeline(Pipeline):
         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
+
+