generate_clusters.py

"""
Run this script to
1. generate clusters by using repository topic embeddings (SciBERT model) + Kmeans
2. generate clusters by using repository code embeddings (UniXCoder model) + KMeans

Files generated by this script:
* repo_topic_clusters.json --- a json file representing repo-topic_cluster dictionary
* kmeans_model_topic_scibert.pkl --- a pickle file for storing kmeans model based on topic embeddings generated by SciBERT model
* repo_code_clusters.json --- a json file representing repo-code_cluster dictionary
* kmeans_model_code_unixcoder.pkl --- a pickle file for storing kmeans model based on code embeddings generated by UniXCoder model
"""

import os
import json
import nltk
import joblib
import numpy
import pandas as pd
import plotly.express as px
import numpy as np
import torch
from pathlib import Path
from common.repo_doc import RepoDoc
from docarray.index import InMemoryExactNNIndex
from transformers import AutoTokenizer, AutoModel
from sklearn.cluster import _kmeans
from sklearn import preprocessing
from sklearn.metrics import pairwise_distances
from nltk.stem import WordNetLemmatizer
from tqdm.auto import tqdm

# from common.skmeans import SKMeans

nltk.download("wordnet")
github_token = os.environ.get("GITHUB_TOKEN")  # Also can be specified by user
INDEX_PATH = Path(__file__).parent.joinpath("index.bin")
device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)

# Initialize SciBERT model and tokenizer
# Option 1 --- Download model by HuggingFace username/model_name
tokenizer = AutoTokenizer.from_pretrained('allenai/scibert_scivocab_uncased')
model = AutoModel.from_pretrained('allenai/scibert_scivocab_uncased').to(device)


# Option 2 --- Download model locally
# tokenizer = AutoTokenizer.from_pretrained("scibert_scivocab_uncased")
# model = AutoModel.from_pretrained('scibert_scivocab_uncased').to(device)


def get_repo_topics():
    """
    The function for getting repositories topics
    :return: repositories names and topics dictionary
    """
    index = InMemoryExactNNIndex[RepoDoc](index_file_path=INDEX_PATH)
    docs = index._docs
    repo_topics = {doc.name: doc.topics for doc in docs}
    return repo_topics


def generate_scibert_embedding(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 = model(**inputs)
    # Use mean pooling for sentence representation
    embeddings = outputs.last_hidden_state.mean(dim=1).cpu().detach().numpy()
    return embeddings


def generate_topic_embeddings(repo_topics):
    """
    The function for generating topic embeddings
    :param repo_topics: repository topics dictionary
    :return: repository topic embeddings dictionary
    """
    lemmatizer = WordNetLemmatizer()
    repo_topic_embeddings = {}
    for repo, topics in tqdm(repo_topics.items()):
        processed_topics = ' '.join([lemmatizer.lemmatize(topic.lower().replace('-', ' ')) for topic in topics if
                                     topic.lower() not in ["python", "python3"]])
        repo_topic_embeddings[repo] = generate_scibert_embedding(processed_topics)[0]

    return repo_topic_embeddings


def generate_topic_clusters(repo_topic_embeddings):
    """
    The function for generating clusters by repository topic embeddings
    :param repo_topic_embeddings: repository topic embeddings dictionary
    :return: repository topic clusters dictionary
    """

    def cosine_dist(X, Y=None):
        return pairwise_distances(X, Y, metric='cosine')

    num_clusters = 150
    vectors = np.array(list(repo_topic_embeddings.values()))

    # Normalised the vectors
    vectors = preprocessing.normalize(vectors)
    # Change the distance calculation function
    _kmeans.euclidean_distances = cosine_dist

    kmeans = _kmeans.KMeans(n_clusters=num_clusters, random_state=0, max_iter=1000).fit(vectors)
    repo_topic_clusters = {repo: int(cluster_label) for repo, cluster_label in
                           zip(repo_topic_embeddings.keys(), kmeans.labels_)}

    # Discard --- the SKMeans model
    # kmeans = SKMeans(no_clusters=num_clusters, iters=1000)
    # kmeans.fit(vectors)
    # repo_topic_clusters = {repo: int(cluster_label) for repo, cluster_label in
    #                        zip(repo_topic_embeddings.keys(), kmeans.labels)}

    joblib.dump(kmeans, "kmeans_model_topic_scibert.pkl")
    with open("repo_topic_clusters.json", "w") as file:
        json.dump(repo_topic_clusters, file, indent=4)

    return repo_topic_clusters


def get_code_embeddings():
    """
    The function to get code embeddings of repositories
    :return: repository code embeddings dictionary
    """
    index = InMemoryExactNNIndex[RepoDoc](index_file_path=INDEX_PATH)
    docs = index._docs
    repo_code_embeddings = {
        doc.name: (doc.code_embedding if doc.code_embedding is not None else numpy.zeros((768,), dtype=np.float32))
        for doc in docs}
    return repo_code_embeddings


def generate_code_clusters(repo_code_embeddings):
    """
    The function for generating clusters by repository code embeddings
    :param repo_code_embeddings: repository code embeddings dictionary
    :return: repository code clusters dictionary
    """

    def cosine_dist(X, Y=None):
        return pairwise_distances(X, Y, metric='cosine')

    num_clusters = 150
    vectors = np.array(list(repo_code_embeddings.values()))

    # Normalised the vectors
    vectors = preprocessing.normalize(vectors)
    # Change the distance calculation function
    _kmeans.euclidean_distances = cosine_dist

    kmeans = _kmeans.KMeans(n_clusters=num_clusters, random_state=0, max_iter=1000).fit(vectors)
    repo_code_clusters = {repo: int(cluster_label) for repo, cluster_label in
                          zip(repo_code_embeddings.keys(), kmeans.labels_)}

    # Discard --- The SKMeans model
    # kmeans = SKMeans(no_clusters=num_clusters, iters=1000)
    # kmeans.fit(vectors)
    # repo_code_clusters = {repo: int(cluster_label) for repo, cluster_label in
    #                       zip(repo_code_embeddings.keys(), kmeans.labels)}

    joblib.dump(kmeans, "kmeans_model_code_unixcoder.pkl")
    with open("repo_code_clusters.json", "w") as file:
        json.dump(repo_code_clusters, file, indent=4)

    return repo_code_clusters


if __name__ == "__main__":
    # # 1. Topic cluster
    # Getting topics for each repository
    repo_topics = get_repo_topics()
    # Generating topic embeddings
    repo_topic_embeddings = generate_topic_embeddings(repo_topics)
    # Generating topic clusters using kmeans model
    repo_topic_clusters = generate_topic_clusters(repo_topic_embeddings)
    # Creating a DataFrame to visualize the repository-topic_cluster assignments
    topic_df = pd.DataFrame(
        {"Repository": list(repo_topic_clusters.keys()), "Topic Cluster": list(repo_topic_clusters.values())})
    # Creating an interactive scatter plot with Plotly
    fig = px.scatter(topic_df, x="Topic Cluster", hover_data=["Repository"],
                     title="Repository Assignments to Topic Clusters with SciBert")
    fig.show()

    # 2. Code cluster
    # Getting code embeddings for each repository
    repo_code_embeddings = get_code_embeddings()
    # Generating code clusters using kmeans model
    repo_code_clusters = generate_code_clusters(repo_code_embeddings)
    # Creating a DataFrame to visualize the repository-code_cluster assignments
    code_df = pd.DataFrame(
        {"Repository": list(repo_code_clusters.keys()), "Code Cluster": list(repo_code_clusters.values())})
    # Creating an interactive scatter plot with Plotly
    fig = px.scatter(code_df, x="Code Cluster", hover_data=["Repository"],
                     title="Repository Assignments to Code Clusters with UniXCoder")
    fig.show()