| import ast |
| import tarfile |
| from ast import AsyncFunctionDef, ClassDef, FunctionDef, Module |
|
|
| import numpy as np |
| import requests |
| import torch |
| from tqdm.auto import tqdm |
| from transformers import Pipeline |
|
|
|
|
| def extract_code_and_docs(text: str): |
| """Extract code and documentation from a Python file. |
| |
| Args: |
| text (str): Source code of a Python file |
| |
| Returns: |
| tuple: A tuple of two sets, the first is the code set, and the second is the docs set, |
| each set contains unique code string or docstring, respectively. |
| """ |
| code_set = set() |
| docs_set = set() |
| root = ast.parse(text) |
| for node in ast.walk(root): |
| if not isinstance(node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)): |
| continue |
| docs = ast.get_docstring(node) |
| node_without_docs = node |
| if docs is not None: |
| docs_set.add(docs) |
| |
| node_without_docs.body = node_without_docs.body[1:] |
| if isinstance(node, (AsyncFunctionDef, FunctionDef)): |
| code_set.add(ast.unparse(node_without_docs)) |
|
|
| return code_set, docs_set |
|
|
|
|
| def get_metadata(repo_name, headers=None): |
| api_url = f"https://api.github.com/repos/{repo_name}" |
| tqdm.write(f"[+] Getting metadata for {repo_name}") |
| try: |
| response = requests.get(api_url, headers=headers) |
| response.raise_for_status() |
|
|
| return response.json() |
| except requests.exceptions.HTTPError as e: |
| tqdm.write(f"[-] Failed to retrieve metadata from {repo_name}: {e}") |
| return {} |
|
|
|
|
| def download_and_extract(repos, headers=None): |
| extracted_infos = [] |
| for repo_name in tqdm(repos, disable=len(repos) <= 1): |
| |
| metadata = get_metadata(repo_name, headers=headers) |
| repo_info = { |
| "name": repo_name, |
| "funcs": set(), |
| "docs": set(), |
| "topics": [], |
| "license": "", |
| "stars": metadata.get("stargazers_count"), |
| } |
| if metadata.get("topics"): |
| repo_info["topics"] = metadata["topics"] |
| if metadata.get("license"): |
| repo_info["license"] = metadata["license"]["spdx_id"] |
|
|
| |
| download_url = f"https://api.github.com/repos/{repo_name}/tarball" |
| tqdm.write(f"[+] Downloading {repo_name}") |
| try: |
| response = requests.get(download_url, headers=headers, stream=True) |
| response.raise_for_status() |
| except requests.exceptions.HTTPError as e: |
| tqdm.write(f"[-] Failed to download {repo_name}: {e}") |
| continue |
|
|
| |
| 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["funcs"].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) |
|
|
| return extracted_infos |
|
|
|
|
| class RepoEmbeddingPipeline(Pipeline): |
| def __init__(self, github_token=None, *args, **kwargs): |
| super().__init__(*args, **kwargs) |
|
|
| self.github_token = github_token |
| if self.github_token: |
| print("[+] GitHub token set!") |
| else: |
| print( |
| "[*] Consider setting GitHub token to avoid hitting rate limits. \n" |
| "For more info, see: " |
| "https://docs.github.com/authentication/keeping-your-account-and-data-secure/creating-a-personal-access-token" |
| ) |
|
|
| def _sanitize_parameters(self, **kwargs): |
| preprocess_kwargs = {} |
| if "github_token" in kwargs: |
| preprocess_kwargs["github_token"] = kwargs["github_token"] |
|
|
| _forward_kwargs = {} |
| if "max_length" in kwargs: |
| _forward_kwargs["max_length"] = kwargs["max_length"] |
| if "st_progress" in kwargs: |
| _forward_kwargs["st_progress"] = kwargs["st_progress"] |
|
|
| return preprocess_kwargs, _forward_kwargs, {} |
|
|
| def preprocess(self, inputs, github_token=None): |
| if isinstance(inputs, str): |
| inputs = [inputs] |
|
|
| headers = {"Accept": "application/vnd.github+json"} |
| token = github_token or self.github_token |
| if token: |
| headers["Authorization"] = f"Bearer {token}" |
|
|
| extracted_infos = download_and_extract(inputs, headers=headers) |
|
|
| return extracted_infos |
|
|
| def encode(self, text, max_length): |
| """ |
| Generates an embedding for a input string. |
| |
| Parameters: |
| |
| * `text`- The input string to be embedded. |
| * `max_length`- The maximum total source sequence length after tokenization. |
| """ |
| assert max_length < 1024 |
|
|
| tokenizer = self.tokenizer |
|
|
| tokens = ( |
| [tokenizer.cls_token, "<encoder-only>", tokenizer.sep_token] |
| + tokenizer.tokenize(text)[: max_length - 4] |
| + [tokenizer.sep_token] |
| ) |
| tokens_id = tokenizer.convert_tokens_to_ids(tokens) |
| source_ids = torch.tensor([tokens_id]).to(self.device) |
|
|
| token_embeddings = self.model(source_ids)[0] |
| sentence_embeddings = token_embeddings.mean(dim=1) |
|
|
| return sentence_embeddings |
|
|
| def _forward(self, extracted_infos, max_length=512, st_progress=None): |
| repo_dataset = [] |
| num_texts = sum(len(x["funcs"]) + len(x["docs"]) for x in extracted_infos) |
| with tqdm(total=num_texts) as pbar: |
| for repo_info in extracted_infos: |
| repo_name = repo_info["name"] |
| entry = { |
| "name": repo_name, |
| "topics": repo_info["topics"], |
| "license": repo_info["license"], |
| "stars": repo_info["stars"], |
| } |
|
|
| pbar.set_description(f"Processing {repo_name}") |
|
|
| tqdm.write(f"[*] Generating embeddings for {repo_name}") |
|
|
| code_embeddings = [] |
| for func in repo_info["funcs"]: |
| code_embeddings.append( |
| [func, self.encode(func, max_length).squeeze().tolist()] |
| ) |
|
|
| pbar.update(1) |
| if st_progress: |
| st_progress.progress(pbar.n / pbar.total) |
|
|
| entry["code_embeddings"] = code_embeddings |
| entry["mean_code_embedding"] = ( |
| np.mean([x[1] for x in code_embeddings], axis=0).tolist() |
| if code_embeddings |
| else None |
| ) |
|
|
| doc_embeddings = [] |
| for doc in repo_info["docs"]: |
| doc_embeddings.append( |
| [doc, self.encode(doc, max_length).squeeze().tolist()] |
| ) |
|
|
| pbar.update(1) |
| if st_progress: |
| st_progress.progress(pbar.n / pbar.total) |
|
|
| entry["doc_embeddings"] = doc_embeddings |
| entry["mean_doc_embedding"] = ( |
| np.mean([x[1] for x in doc_embeddings], axis=0).tolist() |
| if doc_embeddings |
| else None |
| ) |
|
|
| repo_dataset.append(entry) |
|
|
| return repo_dataset |
|
|
| def postprocess(self, repo_dataset): |
| return repo_dataset |
|
|
