Update app.py

app.py

@@ -17,70 +17,21 @@ from PIL import Image
 tokenizer = RobertaTokenizer.from_pretrained("microsoft/graphcodebert-base", cache_dir="models/")
 model = RobertaModel.from_pretrained("microsoft/graphcodebert-base", cache_dir="models/")
 
-# Define sorting algorithms
-sorting_algorithms = {
-    "Bubble_Sort": """
-def bubble_sort(arr):
+# Default sorting algorithm code snippets
+default_code_1 = """def bubble_sort(arr):
     n = len(arr)
     for i in range(n):
         for j in range(0, n-i-1):
             if arr[j] > arr[j+1]:
                 arr[j], arr[j+1] = arr[j+1], arr[j]
-    return arr
-""",
+    return arr"""
 
-    "Selection_Sort": """
-def selection_sort(arr):
-    for i in range(len(arr)):
-        min_idx = i
-        for j in range(i+1, len(arr)):
-            if arr[j] < arr[min_idx]:
-                min_idx = j
-        arr[i], arr[min_idx] = arr[min_idx], arr[i]
-    return arr
-""",
-
-    "Insertion_Sort": """
-def insertion_sort(arr):
-    for i in range(1, len(arr)):
-        key = arr[i]
-        j = i-1
-        while j >= 0 and key < arr[j]:
-            arr[j + 1] = arr[j]
-            j -= 1
-        arr[j + 1] = key
-    return arr
-""",
-
-    "Merge_Sort": """
-def merge_sort(arr):
-    if len(arr) > 1:
-        mid = len(arr) // 2
-        L = arr[:mid]
-        R = arr[mid:]
-        merge_sort(L)
-        merge_sort(R)
-        i = j = k = 0
-        while i < len(L) and j < len(R):
-            if L[i] < R[j]:
-                arr[k] = L[i]
-                i += 1
-            else:
-                arr[k] = R[j]
-                j += 1
-            k += 1
-        while i < len(L):
-            arr[k] = L[i]
-            i += 1
-            k += 1
-        while j < len(R):
-            arr[k] = R[j]
-            j += 1
-            k += 1
-    return arr
-""",
+default_code_2 = """def quick_sort(arr, low, high):
+    if low < high:
+        pi = partition(arr, low, high)
+        quick_sort(arr, low, pi - 1)
+        quick_sort(arr, pi + 1, high)
 
-    "Quick_Sort": """
 def partition(arr, low, high):
     i = (low - 1)
     pivot = arr[high]
@@ -89,15 +40,7 @@ def partition(arr, low, high):
             i += 1
             arr[i], arr[j] = arr[j], arr[i]
     arr[i+1], arr[high] = arr[high], arr[i+1]
-    return (i + 1)
-def quick_sort(arr, low, high):
-    if low < high:
-        pi = partition(arr, low, high)
-        quick_sort(arr, low, pi - 1)
-        quick_sort(arr, pi + 1, high)
-    return arr
-"""
-}
+    return (i + 1)"""
 
 # Get token embeddings for a code snippet
 def get_token_embeddings(code):
@@ -109,10 +52,7 @@ def get_token_embeddings(code):
     return token_embeddings, tokens
 
 # Plot comparison between two algorithms
-def compare_algorithms(algo1_name, algo2_name):
-    code1 = sorting_algorithms[algo1_name]
-    code2 = sorting_algorithms[algo2_name]
-
+def compare_algorithms(code1, code2):
     emb1, tokens1 = get_token_embeddings(code1)
     emb2, tokens2 = get_token_embeddings(code2)
 
@@ -121,8 +61,8 @@ def compare_algorithms(algo1_name, algo2_name):
     coords = pca.fit_transform(combined)
 
     plt.figure(figsize=(6, 5), dpi=150)
-    plt.scatter(coords[:len(tokens1), 0], coords[:len(tokens1), 1], color='red', label=algo1_name, s=20)
-    plt.scatter(coords[len(tokens1):, 0], coords[len(tokens1):, 1], color='blue', label=algo2_name, s=20)
+    plt.scatter(coords[:len(tokens1), 0], coords[:len(tokens1), 1], color='red', label="Code 1", s=20)
+    plt.scatter(coords[len(tokens1):, 0], coords[len(tokens1):, 1], color='blue', label="Code 2", s=20)
     plt.legend()
     plt.xticks([]); plt.yticks([]); plt.grid(False)
 
@@ -136,15 +76,16 @@ def compare_algorithms(algo1_name, algo2_name):
 interface = gr.Interface(
     fn=compare_algorithms,
     inputs=[
-        gr.Dropdown(choices=list(sorting_algorithms.keys()), label="Algorithm 1"),
-        gr.Dropdown(choices=list(sorting_algorithms.keys()), label="Algorithm 2")
+        gr.Textbox(lines=15, label="Code 1", value=default_code_1, language="python"),
+        gr.Textbox(lines=15, label="Code 2", value=default_code_2, language="python")
     ],
     outputs=gr.Image(type="pil", label="Token Embedding PCA"),
     title="GraphCodeBERT Token Embedding Comparison",
-    description="Visual comparison of token-level embeddings from GraphCodeBERT for classical sorting algorithms."
+    description="Edit or paste two Python code snippets. This tool compares their token-level embeddings using GraphCodeBERT and PCA."
 )
 
 if __name__ == "__main__":
     interface.launch()
 
 
+