-
Notifications
You must be signed in to change notification settings - Fork 2
Expand file tree
/
Copy pathRepoSim4Py.py
More file actions
125 lines (106 loc) · 3.82 KB
/
Copy pathRepoSim4Py.py
File metadata and controls
125 lines (106 loc) · 3.82 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
"""
The python file for running this program.
Running example:
>> pip install -r requirements.txt
>> python RepoSim4Py.py -i lepture/authlib idan/oauthlib evonove/django-oauth-toolkit selwin/python-user-agents SmileyChris/django-countries django-compressor/django-compressor billpmurphy/hask pytoolz/toolz Suor/funcy przemyslawjanpietrzak/pyMonet -o output/ --eval
"""
import os
import argparse
import pickle
from itertools import combinations
from pathlib import Path
import pandas as pd
import numpy as np
import torch
from torch.nn.functional import cosine_similarity
from transformers import pipeline
def cossim(embedding1: np.ndarray, embedding2: np.ndarray):
"""
The method for calculating cosine similarity of two numpy array (two vectors).
:param embedding1: array 1, embedding 1
:param embedding2: array 2, embedding 2
:return: the cosine similarity of these two vectors.
"""
embedding1 = torch.tensor(embedding1, dtype=torch.float32).squeeze(0)
embedding2 = torch.tensor(embedding2, dtype=torch.float32).squeeze(0)
return cosine_similarity(embedding1, embedding2, dim=0).item()
def main():
"""
The main method for running this program.
:return: None
"""
# Arguments rules
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument(
"-i",
"--input",
nargs="+",
help="Input repositories",
required=True,
)
parser.add_argument("-o", "--output", help="Output directory", required=True)
parser.add_argument(
"-e",
"--eval",
help="Evaluate cosine similarities between all repository combinations",
action="store_true",
)
args = parser.parse_args()
# Building the output directory
output_dir = Path(args.output)
output_dir.mkdir(exist_ok=True)
# Loading model
model = pipeline(
model="Henry65/RepoSim4Py",
trust_remote_code=True,
device_map="auto",
github_token=os.environ.get("GITHUB_TOKEN")
)
# Getting input
REPOS = args.input
# Generating output by model and saving it
output = model(tuple(REPOS))
with open(output_dir / "output.pkl", "wb") as f:
pickle.dump(output, f)
print(f"[+] Model outputs was saved to {output_dir}/output.pkl!")
# Starting to process evaluation
if not args.eval:
return
if len(REPOS) < 2:
print("[-] At least 2 repositories can be compared for similarity!")
return
# Similarity calculation
rows_list = []
for info1, info2 in combinations(output, 2):
rows_list.append(
{
"repo1": info1["name"],
"repo2": info2["name"],
"topics1": info1["topics"],
"topics2": info2["topics"],
"code_sim": cossim(
info1["mean_code_embedding"], info2["mean_code_embedding"]
),
"doc_sim": cossim(
info1["mean_doc_embedding"], info2["mean_doc_embedding"]
),
"requirement_sim": cossim(
info1["mean_requirement_embedding"], info2["mean_requirement_embedding"]
),
"readme_sim": cossim(
info1["mean_readme_embedding"], info2["mean_readme_embedding"]
),
"repo_sim": cossim(
info1["mean_repo_embedding"], info2["mean_repo_embedding"]
),
}
)
# Saving the calculation result
df = pd.DataFrame(rows_list)
df = df.sort_values("repo_sim", ascending=False).reset_index(drop=True)
df.to_csv(output_dir / "evaluation_result.csv", index=False)
print(f"[+] Evaluation results saved to {output_dir}/evaluation_result.csv!")
if __name__ == "__main__":
main()