RepoSnipy / app.py
Henry65's picture
Update similaritycal model
e0d476d
raw
history blame contribute delete
No virus
22.9 kB
import re
import json
import nltk
import joblib
import torch
import pandas as pd
import numpy as np
import streamlit as st
from pathlib import Path
from docarray import DocList
from docarray.index import InMemoryExactNNIndex
from transformers import pipeline
from transformers import AutoTokenizer, AutoModel
from common.repo_doc import RepoDoc
from nltk.stem import WordNetLemmatizer
from similarityCal.utils import calculate_similarity
nltk.download("wordnet")
KMEANS_TOPIC_MODEL_PATH = Path(__file__).parent.joinpath("data/kmeans_model_topic_scibert.pkl")
KMEANS_CODE_MODEL_PATH = Path(__file__).parent.joinpath("data/kmeans_model_code_unixcoder.pkl")
SCIBERT_MODEL_PATH = "allenai/scibert_scivocab_uncased"
# SCIBERT_MODEL_PATH = Path(__file__).parent.joinpath("data/scibert_scivocab_uncased") # Download locally
device = (
"cuda"
if torch.cuda.is_available()
else "mps"
if torch.backends.mps.is_available()
else "cpu"
)
# 1. Product environment
# INDEX_PATH = Path(__file__).parent.joinpath("data/index.bin")
# TOPIC_CLUSTER_PATH = Path(__file__).parent.joinpath("data/repo_topic_clusters.json")
# CODE_CLUSTER_PATH = Path(__file__).parent.joinpath("data/repo_code_clusters.json")
# 2. Developing environment
INDEX_PATH = Path(__file__).parent.joinpath("data/index_developing.bin")
TOPIC_CLUSTER_PATH = Path(__file__).parent.joinpath("data/repo_topic_clusters_developing.json")
CODE_CLUSTER_PATH = Path(__file__).parent.joinpath("data/repo_code_clusters_developing.json")
@st.cache_resource(show_spinner="Loading repositories basic information...")
def load_index():
"""
The function to load the index file and return a RepoDoc object with default value
:return: index and a RepoDoc object with default value
"""
default_doc = RepoDoc(
name="",
topics=[],
stars=0,
license="",
code_embedding=None,
doc_embedding=None,
readme_embedding=None,
requirement_embedding=None,
repository_embedding=None
)
return InMemoryExactNNIndex[RepoDoc](index_file_path=INDEX_PATH), default_doc
@st.cache_resource(show_spinner="Loading repositories topic clusters...")
def load_repo_topic_clusters():
"""
The function to load the repo-topic_clusters file
:return: a dictionary with the repo-topic_clusters
"""
with open(TOPIC_CLUSTER_PATH, "r") as file:
repo_topic_clusters = json.load(file)
return repo_topic_clusters
@st.cache_resource(show_spinner="Loading repositories code clusters...")
def load_repo_code_clusters():
"""
The function to load the repo-code_clusters file
:return: a dictionary with the repo-code_clusters
"""
with open(CODE_CLUSTER_PATH, "r") as file:
repo_code_clusters = json.load(file)
return repo_code_clusters
@st.cache_resource(show_spinner="Loading RepoSim4Py pipeline model...")
def load_pipeline_model():
"""
The function to load RepoSim4Py pipeline model
:return: a HuggingFace pipeline
"""
# Option 1 --- Download model by HuggingFace username/model_name
model_path = "Henry65/RepoSim4Py"
# Option 2 --- Download model locally
# model_path = Path(__file__).parent.joinpath("data/RepoSim4Py")
return pipeline(
model=model_path,
trust_remote_code=True,
device_map="auto"
)
@st.cache_resource(show_spinner="Loading SciBERT model...")
def load_scibert_model():
"""
The function to load SciBERT model
:return: tokenizer and model
"""
tokenizer = AutoTokenizer.from_pretrained(SCIBERT_MODEL_PATH)
scibert_model = AutoModel.from_pretrained(SCIBERT_MODEL_PATH).to(device)
return tokenizer, scibert_model
@st.cache_resource(show_spinner="Loading KMeans model (topic clusters)...")
def load_topic_kmeans_model():
"""
The function to load KMeans model (topic clusters)
:return: a KMeans model (topic clusters)
"""
return joblib.load(KMEANS_TOPIC_MODEL_PATH)
@st.cache_resource(show_spinner="Loading KMeans model (code clusters)...")
def load_code_kmeans_model():
"""
The function to load KMeans model (code clusters)
:return: a KMeans model (code clusters)
"""
return joblib.load(KMEANS_CODE_MODEL_PATH)
@st.cache_resource(show_spinner="Loading SimilarityCal model...")
def load_similaritycal_model(mode: str):
"""
The function to load SimilarityCal model
mode: 'code' or 'topic'
:return: the SimilarityCal model
"""
if mode == 'topic':
sim_cal_model = torch.load('similarityCal/topic.pt')
elif mode == 'code':
sim_cal_model = torch.load('similarityCal/code.pt')
else:
raise ValueError("parameter 'mode' must be 'code' or 'topic'")
sim_cal_model.to(device)
sim_cal_model.eval()
return sim_cal_model
def generate_scibert_embedding(tokenizer, scibert_model, text):
"""
The function for generating SciBERT embeddings based on topic text
:param text: the topic text
:return: topic embeddings
"""
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
outputs = scibert_model(**inputs)
# Use mean pooling for sentence representation
embeddings = outputs.last_hidden_state.mean(dim=1).cpu().detach().numpy()
return embeddings
@st.cache_data(show_spinner=False)
def run_pipeline_model(_model, repo_name, github_token):
"""
The function to generate repo_info by using pipeline model
:param _model: pipeline
:param repo_name: the name of repository
:param github_token: GitHub token
:return: the information generated by the pipeline
"""
with st.spinner(
f"Downloading and extracting the {repo_name}, this may take a while..."
):
extracted_infos = _model.preprocess(repo_name, github_token=github_token)
if not extracted_infos:
return None
st_proress_bar = st.progress(0.0)
with st.spinner(f"Generating embeddings for {repo_name}..."):
repo_info = _model.forward(extracted_infos, st_progress=st_proress_bar)[0]
st_proress_bar.empty()
return repo_info
def run_index_search(index, query, search_field, limit):
"""
The function to search at index file based on query and limit
:param index: the index
:param query: query
:param search_field: which field to search for
:param limit: page limit
:return: a dataframe with search results
"""
top_matches, scores = index.find(
query=query, search_field=search_field, limit=limit
)
search_results = top_matches.to_dataframe()
search_results["scores"] = scores
return search_results
def run_topic_cluster_search(repo_topic_clusters, repo_name_list):
"""
The function to search topic cluster number for such repositories.
:param repo_topic_clusters: dictionary with repo-topic_clusters
:param repo_name_list: list or array represent repository names
:return: topic cluster number list
"""
topic_clusters = []
for repo_name in repo_name_list:
topic_clusters.append(repo_topic_clusters[repo_name])
return topic_clusters
def run_code_cluster_search(repo_code_clusters, repo_name_list):
"""
The function to search code cluster number for such repositories.
:param repo_code_clusters: dictionary with repo-code_clusters
:param repo_name_list: list or array represent repository names
:return: code cluster number list
"""
code_clusters = []
for repo_name in repo_name_list:
code_clusters.append(repo_code_clusters[repo_name])
return code_clusters
def run_similaritycal_search(index, repo_clusters, model, query_doc, query_cluster_number, limit):
"""
The function to run SimilarityCal model.
:param index: index file
:param repo_clusters: repo-clusters (topic_cluster or code_cluster) json file
:param model: SimilarityCal model
:param query_doc: query repo doc
:param query_cluster_number: query repo cluster number (code or topic)
:param limit: limit
:return: result dataframe
"""
docs = index._docs
result_dl = DocList[RepoDoc]()
e1_list, e2_list = [], []
for doc in docs:
if query_cluster_number != repo_clusters[doc.name]:
continue
if doc.name != query_doc.name:
e1, e2 = (torch.Tensor(query_doc.repository_embedding),
torch.Tensor(doc.repository_embedding))
e1_list.append(e1)
e2_list.append(e2)
result_dl.append(doc)
e1_list = torch.stack(e1_list).to(device)
e2_list = torch.stack(e2_list).to(device)
model.eval()
similarity_scores = calculate_similarity(model, e1_list, e2_list)[:, 1].cpu().detach().numpy()
df = result_dl.to_dataframe()
df["scores"] = similarity_scores
sorted_df = df.sort_values(by='scores', ascending=False).reset_index(drop=True).head(limit)
sorted_df["rankings"] = sorted_df["scores"].rank(ascending=False, method='first').astype(int)
sorted_df.drop(columns="scores", inplace=True)
return sorted_df
if __name__ == "__main__":
# Loading dataset and models
index, default_doc = load_index()
repo_topic_clusters = load_repo_topic_clusters()
repo_code_clusters = load_repo_code_clusters()
pipeline_model = load_pipeline_model()
lemmatizer = WordNetLemmatizer()
tokenizer, scibert_model = load_scibert_model()
topic_kmeans = load_topic_kmeans_model()
code_kmeans = load_code_kmeans_model()
# Setting the sidebar
with st.sidebar:
st.text_input(
label="GitHub Token",
key="github_token",
type="password",
placeholder="Paste your GitHub token here",
help="Consider setting GitHub token to avoid hitting rate limits: https://docs.github.com/authentication/keeping-your-account-and-data-secure/creating-a-personal-access-token",
)
st.slider(
label="Search results limit",
min_value=1,
max_value=100,
value=10,
step=1,
key="search_results_limit",
help="Limit the number of search results",
)
st.multiselect(
label="Display columns",
options=["scores", "name", "topics", "code cluster", "topic cluster", "stars", "license"],
default=["scores", "name", "topics", "code cluster", "topic cluster", "stars", "license"],
help="Select columns to display in the search results",
key="display_columns",
)
# Setting the main content
st.title("RepoSnipy")
st.text_input(
"Enter a GitHub repository URL or owner/repository (case-sensitive):",
value="",
max_chars=200,
placeholder="numpy/numpy",
key="repo_input",
)
st.checkbox(
label="Add/Update this repository to the index",
value=False,
key="update_index",
help="Encode the latest version of this repository and add/update it to the index",
)
# Setting the search button
search = st.button("Search")
# The regular expression for repository
repo_regex = r"^((git@|http(s)?://)?(github\.com)(/|:))?(?P<owner>[\w.-]+)(/)(?P<repo>[\w.-]+?)(\.git)?(/)?$"
if search:
match_res = re.match(repo_regex, st.session_state.repo_input)
# 1. Repository can be matched
if match_res is not None:
repo_name = f"{match_res.group('owner')}/{match_res.group('repo')}"
records = index.filter({"name": {"$eq": repo_name}})
# 1) Building the query information
query_doc = default_doc.copy() if not records else records[0]
# 2) Recording the topic and code cluster numbers
topic_cluster_number = -1 if not records else repo_topic_clusters[repo_name]
code_cluster_number = -1 if not records else repo_code_clusters[repo_name]
# Importance 1 ---- situation need to update repository information and cluster numbers
if st.session_state.update_index or not records:
# 1) Updating repository information by using RepoSim4Py pipeline
repo_info = run_pipeline_model(pipeline_model, repo_name, st.session_state.github_token)
if repo_info is None:
st.error("Repository not found or invalid GitHub token!")
st.stop()
query_doc.name = repo_info["name"]
query_doc.topics = repo_info["topics"]
query_doc.stars = repo_info["stars"]
query_doc.license = repo_info["license"]
query_doc.code_embedding = None if np.all(repo_info["mean_code_embedding"] == 0) else repo_info[
"mean_code_embedding"].reshape(-1)
query_doc.doc_embedding = None if np.all(repo_info["mean_doc_embedding"] == 0) else repo_info[
"mean_doc_embedding"].reshape(-1)
query_doc.readme_embedding = None if np.all(repo_info["mean_readme_embedding"] == 0) else repo_info[
"mean_readme_embedding"].reshape(-1)
query_doc.requirement_embedding = None if np.all(repo_info["mean_requirement_embedding"] == 0) else \
repo_info["mean_requirement_embedding"].reshape(-1)
query_doc.repository_embedding = None if np.all(repo_info["mean_repo_embedding"] == 0) else repo_info[
"mean_repo_embedding"].reshape(-1)
# 2) Updating topic cluster number
topics_text = ' '.join(
[lemmatizer.lemmatize(topic.lower().replace('-', ' ')) for topic in query_doc.topics if
topic.lower() not in ["python", "python3"]])
topic_embeddings = generate_scibert_embedding(tokenizer, scibert_model, topics_text)
topic_cluster_number = int(topic_kmeans.predict(topic_embeddings)[0])
# 3) Updating code cluster number
code_embeddings = np.zeros((768,),
dtype=np.float32) if query_doc.code_embedding is None else query_doc.code_embedding
code_cluster_number = int(code_kmeans.predict(code_embeddings.reshape(1, -1))[0])
# Importance 2 ---- update index file and repository clusters (topic and code) files
if st.session_state.update_index:
if not query_doc.license:
st.warning(
"License is missing in this repository and will not be persisted!"
)
elif (query_doc.code_embedding is None) and (query_doc.doc_embedding is None) and (
query_doc.requirement_embedding is None) and (query_doc.readme_embedding is None) and (
query_doc.repository_embedding is None):
st.warning(
"This repository has no such useful information (code, docstring, readme and requirement) extracted and will not be persisted!"
)
else:
index.index(query_doc)
repo_topic_clusters[query_doc.name] = topic_cluster_number
repo_code_clusters[query_doc.name] = code_cluster_number
with st.spinner("Persisting the index and repository clusters (topic and code)..."):
index.persist(str(INDEX_PATH))
with open(TOPIC_CLUSTER_PATH, "w") as file:
json.dump(repo_topic_clusters, file, indent=4)
with open(CODE_CLUSTER_PATH, "w") as file:
json.dump(repo_code_clusters, file, indent=4)
st.success("Repository updated to the index!")
load_index.clear()
load_repo_topic_clusters.clear()
load_repo_code_clusters.clear()
st.session_state["query_doc"] = query_doc
st.session_state["topic_cluster_number"] = topic_cluster_number
st.session_state["code_cluster_number"] = code_cluster_number
# 2. Repository cannot be matched
else:
st.error("Invalid input!")
# Starting to query
if "query_doc" in st.session_state:
query_doc = st.session_state.query_doc
topic_cluster_number = st.session_state.topic_cluster_number
code_cluster_number = st.session_state.code_cluster_number
limit = st.session_state.search_results_limit
# Showing the query repository information
st.dataframe(
pd.DataFrame(
[
{
"name": query_doc.name,
"topics": query_doc.topics,
"code cluster": code_cluster_number,
"topic cluster": topic_cluster_number,
"stars": query_doc.stars,
"license": query_doc.license,
}
],
)
)
display_columns = st.session_state.display_columns
modified_display_columns = ["rankings" if col == "scores" else col for col in display_columns]
code_sim_tab, doc_sim_tab, readme_sim_tab, requirement_sim_tab, repo_sim_tab, code_cluster_tab, topic_cluster_tab, = st.tabs(
["Code_sim", "Docstring_sim", "Readme_sim", "Requirement_sim",
"Repository_sim", "Code_cluster_sim", "Topic_cluster_sim"])
with code_sim_tab:
if query_doc.code_embedding is not None:
code_sim_res = run_index_search(index, query_doc, "code_embedding", limit)
topic_cluster_numbers = run_topic_cluster_search(repo_topic_clusters, code_sim_res["name"])
code_sim_res["topic cluster"] = topic_cluster_numbers
code_cluster_numbers = run_code_cluster_search(repo_code_clusters, code_sim_res["name"])
code_sim_res["code cluster"] = code_cluster_numbers
st.dataframe(code_sim_res[display_columns])
else:
st.error("No function code was extracted for this repository!")
with doc_sim_tab:
if query_doc.doc_embedding is not None:
doc_sim_res = run_index_search(index, query_doc, "doc_embedding", limit)
topic_cluster_numbers = run_topic_cluster_search(repo_topic_clusters, doc_sim_res["name"])
doc_sim_res["topic cluster"] = topic_cluster_numbers
code_cluster_numbers = run_code_cluster_search(repo_code_clusters, doc_sim_res["name"])
doc_sim_res["code cluster"] = code_cluster_numbers
st.dataframe(doc_sim_res[display_columns])
else:
st.error("No function docstring was extracted for this repository!")
with readme_sim_tab:
if query_doc.readme_embedding is not None:
readme_sim_res = run_index_search(index, query_doc, "readme_embedding", limit)
topic_cluster_numbers = run_topic_cluster_search(repo_topic_clusters, readme_sim_res["name"])
readme_sim_res["topic cluster"] = topic_cluster_numbers
code_cluster_numbers = run_code_cluster_search(repo_code_clusters, readme_sim_res["name"])
readme_sim_res["code cluster"] = code_cluster_numbers
st.dataframe(readme_sim_res[display_columns])
else:
st.error("No readme file was extracted for this repository!")
with requirement_sim_tab:
if query_doc.requirement_embedding is not None:
requirement_sim_res = run_index_search(index, query_doc, "requirement_embedding", limit)
topic_cluster_numbers = run_topic_cluster_search(repo_topic_clusters, requirement_sim_res["name"])
requirement_sim_res["topic cluster"] = topic_cluster_numbers
code_cluster_numbers = run_code_cluster_search(repo_code_clusters, requirement_sim_res["name"])
requirement_sim_res["code cluster"] = code_cluster_numbers
st.dataframe(requirement_sim_res[display_columns])
else:
st.error("No requirement file was extracted for this repository!")
with repo_sim_tab:
if query_doc.repository_embedding is not None:
# Repo Sim tab
repo_sim_res = run_index_search(index, query_doc, "repository_embedding", limit)
topic_cluster_numbers = run_topic_cluster_search(repo_topic_clusters, repo_sim_res["name"])
repo_sim_res["topic cluster"] = topic_cluster_numbers
code_cluster_numbers = run_code_cluster_search(repo_code_clusters, repo_sim_res["name"])
repo_sim_res["code cluster"] = code_cluster_numbers
st.dataframe(repo_sim_res[display_columns])
else:
st.error("No such useful information was extracted for this repository!")
with code_cluster_tab:
if query_doc.repository_embedding is not None:
sim_cal_model = load_similaritycal_model("code")
cluster_df = run_similaritycal_search(index, repo_code_clusters, sim_cal_model,
query_doc, code_cluster_number, limit)
code_cluster_numbers = run_code_cluster_search(repo_code_clusters, cluster_df["name"])
cluster_df["code cluster"] = code_cluster_numbers
topic_cluster_numbers = run_topic_cluster_search(repo_topic_clusters, cluster_df["name"])
cluster_df["topic cluster"] = topic_cluster_numbers
st.dataframe(cluster_df[modified_display_columns])
else:
st.error("No such useful information was extracted for this repository!")
with topic_cluster_tab:
if query_doc.repository_embedding is not None:
sim_cal_model = load_similaritycal_model("topic")
cluster_df = run_similaritycal_search(index, repo_topic_clusters, sim_cal_model,
query_doc, topic_cluster_number, limit)
topic_cluster_numbers = run_topic_cluster_search(repo_topic_clusters, cluster_df["name"])
cluster_df["topic cluster"] = topic_cluster_numbers
code_cluster_numbers = run_code_cluster_search(repo_code_clusters, cluster_df["name"])
cluster_df["code cluster"] = code_cluster_numbers
st.dataframe(cluster_df[modified_display_columns])
else:
topic_cluster_tab.error("No such useful information was extracted for this repository!")