Update app.py

app.py

@@ -2,7 +2,6 @@ import javalang
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch_geometric.data import Data
 import re
 import gradio as gr
 from transformers import AutoTokenizer, AutoModel
@@ -10,186 +9,61 @@ from pathlib import Path
 
 # Configuration
 MAX_FILE_SIZE = 5000
-MAX_AST_DEPTH = 50
 EMBEDDING_DIM = 128
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# AST Encoder
-class ASTNodeEncoder:
-    def __init__(self):
-        self.node_types = set()
-        self.type_to_idx = {}
-        
-    def fit(self, ast_nodes):
-        self.node_types.update(ast_nodes)
-        self.type_to_idx = {t: i for i, t in enumerate(sorted(self.node_types))}
-        
-    def encode(self, node_type):
-        if node_type not in self.type_to_idx:
-            return torch.zeros(EMBEDDING_DIM)
-        idx = self.type_to_idx[node_type]
-        embedding = torch.zeros(EMBEDDING_DIM)
-        embedding[idx % EMBEDDING_DIM] = 1
-        embedding += torch.randn(EMBEDDING_DIM) * 0.1
-        return embedding
-
-# Code Normalization
-def normalize_java_code(code):
-    code = re.sub(r'//.*', '', code)
-    code = re.sub(r'/\*.*?\*/', '', code, flags=re.DOTALL)
-    code = re.sub(r'"[^"]*"', '"<STRING>"', code)
-    code = re.sub(r"'[^']*'", "'<CHAR>'", code)
-    return ' '.join(code.split())
-
-# AST Processing
-def extract_ast_paths(node, encoder, current_path=None, paths=None, depth=0):
-    if current_path is None:
-        current_path = []
-    if paths is None:
-        paths = []
-    
-    if depth > MAX_AST_DEPTH:
-        return paths
-    
-    node_type = str(type(node).__name__)
-    node_embedding = encoder.encode(node_type)
-    current_path.append(node_embedding)
-    
-    if not hasattr(node, 'children') or depth == MAX_AST_DEPTH:
-        paths.append(torch.stack(current_path))
-        current_path.pop()
-        return paths
-    
-    for child in node.children:
-        if isinstance(child, (javalang.ast.Node, list, tuple)):
-            if isinstance(child, (list, tuple)):
-                for c in child:
-                    if isinstance(c, javalang.ast.Node):
-                        extract_ast_paths(c, encoder, current_path, paths, depth+1)
-            elif isinstance(child, javalang.ast.Node):
-                extract_ast_paths(child, encoder, current_path, paths, depth+1)
-    
-    current_path.pop()
-    return paths
-
-def ast_to_graph_data(ast, encoder):
-    paths = extract_ast_paths(ast, encoder)
-    if not paths:
-        return None
-    
-    edge_index = []
-    node_features = []
-    node_counter = 0
-    node_mapping = {}
-    
-    for path in paths:
-        for i in range(len(path) - 1):
-            for j in [i, i+1]:
-                node_key = tuple(path[j].tolist())
-                if node_key not in node_mapping:
-                    node_mapping[node_key] = node_counter
-                    node_features.append(path[j])
-                    node_counter += 1
-            
-            src = node_mapping[tuple(path[i].tolist())]
-            dst = node_mapping[tuple(path[i+1].tolist())]
-            edge_index.append([src, dst])
-    
-    if not edge_index:
-        return None
-    
-    node_features = torch.stack(node_features)
-    edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()
-    return Data(x=node_features, edge_index=edge_index)
-
-# Model Architecture
-class ASTGNN(nn.Module):
-    def __init__(self, input_dim, hidden_dim):
-        super().__init__()
-        self.conv1 = nn.Sequential(
-            nn.Linear(input_dim, hidden_dim),
-            nn.ReLU(),
-            nn.Linear(hidden_dim, hidden_dim)
-        )
-        self.conv2 = nn.Sequential(
-            nn.Linear(hidden_dim, hidden_dim),
-            nn.ReLU(),
-            nn.Linear(hidden_dim, hidden_dim)
-        )
-        self.pool = nn.AdaptiveMaxPool1d(1)
-        
-    def forward(self, data):
-        x, edge_index = data.x.to(DEVICE), data.edge_index.to(DEVICE)
-        x = self.conv1(x)
-        x = self.conv2(x)
-        x = x.t().unsqueeze(0)
-        x = self.pool(x)
-        return x.squeeze(0).squeeze(-1)
+# Initialize models once at startup
+tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
+code_model = AutoModel.from_pretrained("microsoft/codebert-base").to(DEVICE)
 
-class HybridCloneDetector(nn.Module):
-    def __init__(self, ast_input_dim, hidden_dim):
+# Simplified model architecture
+class CloneDetector(nn.Module):
+    def __init__(self, hidden_dim):
         super().__init__()
-        self.ast_gnn = ASTGNN(ast_input_dim, hidden_dim)
         self.classifier = nn.Sequential(
             nn.Linear(hidden_dim * 2, hidden_dim),
             nn.ReLU(),
             nn.Linear(hidden_dim, 2))
         
-    def forward(self, ast_data, code_embedding):
-        ast_embed = self.ast_gnn(ast_data)
-        combined = torch.cat([ast_embed, code_embedding.squeeze(0)], dim=0)
-        return self.classifier(combined.unsqueeze(0))
-
-# Load Models
-def load_models():
-    ast_encoder = ASTNodeEncoder()
-    ast_encoder.fit(['MethodDeclaration', 'VariableDeclaration', 'IfStatement'])
-    
-    tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
-    code_model = AutoModel.from_pretrained("microsoft/codebert-base").to(DEVICE)
-    
-    model = HybridCloneDetector(EMBEDDING_DIM, EMBEDDING_DIM).to(DEVICE)
-    if Path('model.pth').exists():
-        model.load_state_dict(torch.load('model.pth', map_location=DEVICE))
-    
-    return ast_encoder, tokenizer, code_model, model
+    def forward(self, emb1, emb2):
+        combined = torch.cat([emb1, emb2], dim=-1)
+        return self.classifier(combined)
 
-ast_encoder, tokenizer, code_model, model = load_models()
+model = CloneDetector(768).to(DEVICE)  # 768 is CodeBERT's hidden size
 
-# Prediction Function
-def predict_clone(code1, code2):
+def get_code_embedding(code):
+    """Get embedding for a single code snippet"""
     try:
-        # Process first code
-        norm_code1 = normalize_java_code(code1)
-        tokens1 = list(javalang.tokenizer.tokenize(norm_code1))
-        parser = javalang.parser.Parser(tokens1)
-        ast1 = parser.parse()
-        ast_data1 = ast_to_graph_data(ast1, ast_encoder)
+        # Normalize code
+        code = re.sub(r'//.*', '', code)
+        code = re.sub(r'/\*.*?\*/', '', code, flags=re.DOTALL)
+        code = ' '.join(code.split())
         
-        inputs1 = tokenizer(norm_code1, return_tensors="pt", truncation=True).to(DEVICE)
+        # Tokenize and get embedding
+        inputs = tokenizer(code, return_tensors="pt", truncation=True, max_length=512).to(DEVICE)
         with torch.no_grad():
-            code_embed1 = code_model(**inputs1).last_hidden_state.mean(dim=1)
-        
-        # Process second code
-        norm_code2 = normalize_java_code(code2)
-        tokens2 = list(javalang.tokenizer.tokenize(norm_code2))
-        parser = javalang.parser.Parser(tokens2)
-        ast2 = parser.parse()
-        ast_data2 = ast_to_graph_data(ast2, ast_encoder)
-        
-        inputs2 = tokenizer(norm_code2, return_tensors="pt", truncation=True).to(DEVICE)
-        with torch.no_grad():
-            code_embed2 = code_model(**inputs2).last_hidden_state.mean(dim=1)
+            outputs = code_model(**inputs)
+        return outputs.last_hidden_state.mean(dim=1)  # Pooled representation
+    except Exception:
+        return torch.zeros(1, 768).to(DEVICE)
+
+def predict_clone(code1, code2):
+    """Compare two code snippets"""
+    try:
+        # Get embeddings
+        emb1 = get_code_embedding(code1)
+        emb2 = get_code_embedding(code2)
         
-        # Predict
+        # Calculate similarity
         with torch.no_grad():
-            logits1 = model(ast_data1.to(DEVICE), code_embed1)
-            logits2 = model(ast_data2.to(DEVICE), code_embed2)
-            sim_score = F.cosine_similarity(logits1, logits2).item()
+            sim_score = F.cosine_similarity(emb1, emb2).item()
+            logits = model(emb1, emb2)
+            prob = F.softmax(logits, dim=-1)[0, 1].item()
         
         return {
-            "Similarity": f"{sim_score:.3f}",
-            "Clone": "Yes" if sim_score > 0.7 else "No"
+            "Similarity Score": f"{sim_score:.3f}",
+            "Clone Probability": f"{prob:.3f}",
+            "Prediction": "Clone" if prob > 0.5 else "Not Clone"
         }
     except Exception as e:
         return {"Error": str(e)}
@@ -201,7 +75,7 @@ demo = gr.Interface(
         gr.Textbox(label="First Java Code", lines=10),
         gr.Textbox(label="Second Java Code", lines=10)
     ],
-    outputs=gr.JSON(label="Prediction"),
+    outputs=gr.JSON(label="Results"),
     examples=[
         ["""public class Hello {
     public static void main(String[] args) {
@@ -221,7 +95,7 @@ demo = gr.Interface(
 }"""]
     ],
     title="Java Code Clone Detector",
-    description="Detect code clones between two Java code snippets using AST and neural embeddings"
+    description="Compare two Java code snippets to detect potential clones"
 )
 
 if __name__ == "__main__":