jorgemarcc/graphcodebert-interpretability

1# -*- coding: utf-8 -*-2"""3Martinez-Gil, J. (2025). Augmenting the Interpretability of GraphCodeBERT for Code Similarity Tasks.4International Journal of Software Engineering and Knowledge Engineering, 35(05), 657–678.5"""6 7import numpy as np8import matplotlib.pyplot as plt9from sklearn.decomposition import PCA10from transformers import RobertaTokenizer, RobertaModel11import torch12import gradio as gr13from io import BytesIO14from PIL import Image15 16# Load GraphCodeBERT from Hugging Face (with cache)17tokenizer = RobertaTokenizer.from_pretrained("microsoft/graphcodebert-base", cache_dir="models/")18model = RobertaModel.from_pretrained("microsoft/graphcodebert-base", cache_dir="models/")19 20# Default sorting algorithm code snippets21default_code_1 = """def bubble_sort(arr):22    n = len(arr)23    for i in range(n):24        for j in range(0, n-i-1):25            if arr[j] > arr[j+1]:26                arr[j], arr[j+1] = arr[j+1], arr[j]27    return arr"""28 29default_code_2 = """def quick_sort(arr, low, high):30    if low < high:31        pi = partition(arr, low, high)32        quick_sort(arr, low, pi - 1)33        quick_sort(arr, pi + 1, high)34 35def partition(arr, low, high):36    i = (low - 1)37    pivot = arr[high]38    for j in range(low, high):39        if arr[j] <= pivot:40            i += 141            arr[i], arr[j] = arr[j], arr[i]42    arr[i+1], arr[high] = arr[high], arr[i+1]43    return (i + 1)"""44 45# Get token embeddings for a code snippet46def get_token_embeddings(code):47    inputs = tokenizer(code, return_tensors="pt", max_length=512, truncation=True, padding=True)48    with torch.no_grad():49        outputs = model(**inputs)50    token_embeddings = outputs.last_hidden_state.squeeze(0).cpu().numpy()51    tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'].squeeze())52    return token_embeddings, tokens53 54# Plot comparison between two algorithms55def compare_algorithms(code1, code2):56    emb1, tokens1 = get_token_embeddings(code1)57    emb2, tokens2 = get_token_embeddings(code2)58 59    combined = np.concatenate([emb1, emb2], axis=0)60    pca = PCA(n_components=2)61    coords = pca.fit_transform(combined)62 63    plt.figure(figsize=(6, 5), dpi=150)64    plt.scatter(coords[:len(tokens1), 0], coords[:len(tokens1), 1], color='red', label="Code 1", s=20)65    plt.scatter(coords[len(tokens1):, 0], coords[len(tokens1):, 1], color='blue', label="Code 2", s=20)66    plt.legend()67    plt.xticks([]); plt.yticks([]); plt.grid(False)68 69    buf = BytesIO()70    plt.savefig(buf, format='png', bbox_inches='tight')71    plt.close()72    buf.seek(0)73    return Image.open(buf)74 75interface = gr.Interface(76    fn=compare_algorithms,77    inputs=[78        gr.Code(language="python", value=default_code_1, label="Code 1"),79        gr.Code(language="python", value=default_code_2, label="Code 2")80    ],81    outputs=gr.Image(type="pil", label="Token Embedding PCA"),82    title="GraphCodeBERT Token Embedding Comparison",83    description="Edit or paste two Python code snippets. This tool compares their token-level embeddings using GraphCodeBERT and PCA.",84    article="""85**Citation**  86Martinez-Gil, J. (2025). *Augmenting the Interpretability of GraphCodeBERT for Code Similarity Tasks.* International Journal of Software Engineering and Knowledge Engineering, 35(05), 657–678.87 88**GitHub Repository**  89[View Source on GitHub](https://github.com/jorge-martinez-gil/graphcodebert-interpretability)90"""91)92 93if __name__ == "__main__":94    interface.launch()95 96 97 98