Create app.py

app.py

@@ -0,0 +1,520 @@
+import streamlit as st
+import javalang
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import re
+import numpy as np
+import networkx as nx
+from transformers import AutoTokenizer, AutoModel
+from torch_geometric.data import Data
+from torch_geometric.nn import GCNConv
+import warnings
+import pandas as pd
+import zipfile
+import os
+from collections import defaultdict
+
+# Set up page config
+st.set_page_config(
+    page_title="Advanced Java Code Clone Detector (IJaDataset 2.1)",
+    page_icon="🔍",
+    layout="wide"
+)
+
+# Suppress warnings
+warnings.filterwarnings("ignore")
+
+# Constants
+MODEL_NAME = "microsoft/codebert-base"
+MAX_LENGTH = 512
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+DATASET_PATH = "archive (1).zip"  # Update this path if needed
+
+# Initialize models with caching
+@st.cache_resource
+def load_models():
+    try:
+        # Load CodeBERT for semantic analysis
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        code_model = AutoModel.from_pretrained(MODEL_NAME).to(DEVICE)
+        
+        # Initialize RNN model
+        class RNNModel(nn.Module):
+            def __init__(self, input_size, hidden_size, num_layers):
+                super(RNNModel, self).__init__()
+                self.hidden_size = hidden_size
+                self.num_layers = num_layers
+                self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True)
+                self.fc = nn.Linear(hidden_size, 1)
+                
+            def forward(self, x):
+                h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(DEVICE)
+                out, _ = self.rnn(x, h0)
+                out = self.fc(out[:, -1, :])
+                return out
+        
+        rnn_model = RNNModel(input_size=768, hidden_size=256, num_layers=2).to(DEVICE)
+        
+        # Initialize GNN model
+        class GNNModel(nn.Module):
+            def __init__(self, node_features):
+                super(GNNModel, self).__init__()
+                self.conv1 = GCNConv(node_features, 128)
+                self.conv2 = GCNConv(128, 64)
+                self.fc = nn.Linear(64, 1)
+                
+            def forward(self, data):
+                x, edge_index = data.x, data.edge_index
+                x = F.relu(self.conv1(x, edge_index))
+                x = F.dropout(x, training=self.training)
+                x = self.conv2(x, edge_index)
+                x = self.fc(x)
+                return torch.sigmoid(x.mean())
+        
+        gnn_model = GNNModel(node_features=128).to(DEVICE)
+        
+        return tokenizer, code_model, rnn_model, gnn_model
+    except Exception as e:
+        st.error(f"Failed to load models: {str(e)}")
+        return None, None, None, None
+
+@st.cache_resource
+def load_dataset():
+    try:
+        # Extract dataset if needed
+        if not os.path.exists("Diverse_100K_Dataset"):
+            with zipfile.ZipFile(DATASET_PATH, 'r') as zip_ref:
+                zip_ref.extractall(".")
+        
+        # Load sample pairs (modify this based on your dataset structure)
+        clone_pairs = []
+        base_path = "Subject_CloneTypes_Directories"
+        
+        # Load pairs from all clone types
+        for clone_type in ["Clone_Type1", "Clone_Type2", "Clone_Type3 - ST", "Clone_Type4"]:
+            type_path = os.path.join(base_path, clone_type)
+            if os.path.exists(type_path):
+                for root, _, files in os.walk(type_path):
+                    if files:
+                        # Take first two files as a pair
+                        if len(files) >= 2:
+                            with open(os.path.join(root, files[0]), 'r', encoding='utf-8') as f1:
+                                code1 = f1.read()
+                            with open(os.path.join(root, files[1]), 'r', encoding='utf-8') as f2:
+                                code2 = f2.read()
+                            clone_pairs.append({
+                                "type": clone_type,
+                                "code1": code1,
+                                "code2": code2
+                            })
+                        break  # Just take one pair per type for demo
+        
+        return clone_pairs[:10]  # Return first 10 pairs for demo
+        
+    except Exception as e:
+        st.error(f"Error loading dataset: {str(e)}")
+        return []
+
+tokenizer, code_model, rnn_model, gnn_model = load_models()
+dataset_pairs = load_dataset()
+
+# AST Processing Functions
+def parse_ast(code):
+    try:
+        tokens = javalang.tokenizer.tokenize(code)
+        parser = javalang.parser.Parser(tokens)
+        tree = parser.parse()
+        return tree
+    except Exception as e:
+        st.warning(f"AST parsing error: {str(e)}")
+        return None
+
+def build_ast_graph(ast_tree):
+    if not ast_tree:
+        return None
+    
+    G = nx.DiGraph()
+    node_id = 0
+    node_map = {}
+    
+    def traverse(node, parent_id=None):
+        nonlocal node_id
+        current_id = node_id
+        node_label = str(type(node).__name__)
+        node_map[current_id] = {'type': node_label, 'node': node}
+        G.add_node(current_id, type=node_label)
+        
+        if parent_id is not None:
+            G.add_edge(parent_id, current_id)
+        
+        node_id += 1
+        
+        for child in node.children:
+            if isinstance(child, javalang.ast.Node):
+                traverse(child, current_id)
+            elif isinstance(child, (list, tuple)):
+                for item in child:
+                    if isinstance(item, javalang.ast.Node):
+                        traverse(item, current_id)
+    
+    traverse(ast_tree)
+    return G, node_map
+
+def ast_to_pyg_data(ast_graph):
+    if not ast_graph:
+        return None
+    
+    # Convert AST to PyTorch Geometric Data format
+    node_features = []
+    node_types = []
+    
+    for node in ast_graph.nodes():
+        node_type = ast_graph.nodes[node]['type']
+        node_types.append(node_type)
+        # Simple one-hot encoding of node types (in practice, use better encoding)
+        feature = [0] * 50  # Assuming max 50 node types
+        feature[hash(node_type) % 50] = 1
+        node_features.append(feature)
+    
+    # Convert networkx graph to edge_index format
+    edge_index = list(ast_graph.edges())
+    if not edge_index:
+        # Add self-loop if no edges
+        edge_index = [(0, 0)]
+    
+    edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()
+    x = torch.tensor(node_features, dtype=torch.float)
+    
+    return Data(x=x, edge_index=edge_index)
+
+# Normalization function
+def normalize_code(code):
+    try:
+        code = re.sub(r'//.*', '', code)  # Remove single-line comments
+        code = re.sub(r'/\*.*?\*/', '', code, flags=re.DOTALL)  # Multi-line comments
+        code = re.sub(r'\s+', ' ', code).strip()  # Normalize whitespace
+        return code
+    except Exception:
+        return code
+
+# Feature extraction functions
+def get_lexical_features(code):
+    """Extract lexical features (for Type-1 and Type-2 clones)"""
+    normalized = normalize_code(code)
+    tokens = re.findall(r'\b\w+\b', normalized)
+    return {
+        'token_count': len(tokens),
+        'unique_tokens': len(set(tokens)),
+        'avg_token_length': np.mean([len(t) for t in tokens]) if tokens else 0
+    }
+
+def get_syntactic_features(ast_tree):
+    """Extract syntactic features (for Type-3 clones)"""
+    if not ast_tree:
+        return {}
+    
+    # Count different node types in AST
+    node_counts = defaultdict(int)
+    
+    def count_nodes(node):
+        node_counts[type(node).__name__] += 1
+        for child in node.children:
+            if isinstance(child, javalang.ast.Node):
+                count_nodes(child)
+            elif isinstance(child, (list, tuple)):
+                for item in child:
+                    if isinstance(item, javalang.ast.Node):
+                        count_nodes(item)
+    
+    count_nodes(ast_tree)
+    return dict(node_counts)
+
+def get_semantic_features(code):
+    """Extract semantic features (for Type-4 clones)"""
+    embedding = get_embedding(code)
+    return embedding.cpu().numpy().flatten() if embedding is not None else None
+
+# Embedding generation
+def get_embedding(code):
+    try:
+        code = normalize_code(code)
+        inputs = tokenizer(
+            code,
+            return_tensors="pt",
+            truncation=True,
+            max_length=MAX_LENGTH,
+            padding='max_length'
+        ).to(DEVICE)
+        
+        with torch.no_grad():
+            outputs = code_model(**inputs)
+        
+        return outputs.last_hidden_state.mean(dim=1)  # Pooled embedding
+    except Exception as e:
+        st.error(f"Error processing code: {str(e)}")
+        return None
+
+# Clone detection models
+def rnn_similarity(emb1, emb2):
+    """Calculate similarity using RNN model"""
+    if emb1 is None or emb2 is None:
+        return None
+    
+    # Prepare input for RNN (sequence of embeddings)
+    combined = torch.cat([emb1.unsqueeze(0), emb2.unsqueeze(0)], dim=0)
+    with torch.no_grad():
+        similarity = rnn_model(combined.permute(1, 0, 2))
+    return torch.sigmoid(similarity).item()
+
+def gnn_similarity(ast1, ast2):
+    """Calculate similarity using GNN model"""
+    if ast1 is None or ast2 is None:
+        return None
+    
+    data1 = ast_to_pyg_data(ast1)
+    data2 = ast_to_pyg_data(ast2)
+    
+    if data1 is None or data2 is None:
+        return None
+    
+    # Move data to device
+    data1 = data1.to(DEVICE)
+    data2 = data2.to(DEVICE)
+    
+    with torch.no_grad():
+        sim1 = gnn_model(data1)
+        sim2 = gnn_model(data2)
+    
+    return F.cosine_similarity(sim1, sim2).item()
+
+def hybrid_similarity(code1, code2):
+    """Combined similarity score using all models"""
+    # Get embeddings
+    emb1 = get_embedding(code1)
+    emb2 = get_embedding(code2)
+    
+    # Parse ASTs
+    ast_tree1 = parse_ast(code1)
+    ast_tree2 = parse_ast(code2)
+    
+    ast_graph1 = build_ast_graph(ast_tree1) if ast_tree1 else None
+    ast_graph2 = build_ast_graph(ast_tree2) if ast_tree2 else None
+    
+    # Calculate individual similarities
+    codebert_sim = F.cosine_similarity(emb1, emb2).item() if emb1 is not None and emb2 is not None else 0
+    rnn_sim = rnn_similarity(emb1, emb2) if emb1 is not None and emb2 is not None else 0
+    gnn_sim = gnn_similarity(ast_graph1[0] if ast_graph1 else None, 
+                            ast_graph2[0] if ast_graph2 else None) or 0
+    
+    # Combine with weights (can be tuned)
+    weights = {
+        'codebert': 0.4,
+        'rnn': 0.3,
+        'gnn': 0.3
+    }
+    
+    combined = (weights['codebert'] * codebert_sim + 
+               weights['rnn'] * rnn_sim + 
+               weights['gnn'] * gnn_sim)
+    
+    return {
+        'combined': combined,
+        'codebert': codebert_sim,
+        'rnn': rnn_sim,
+        'gnn': gnn_sim
+    }
+
+# Comparison function
+def compare_code(code1, code2):
+    if not code1 or not code2:
+        return None
+    
+    with st.spinner('Analyzing code with multiple techniques...'):
+        # Get lexical features
+        lex1 = get_lexical_features(code1)
+        lex2 = get_lexical_features(code2)
+        
+        # Get AST trees
+        ast_tree1 = parse_ast(code1)
+        ast_tree2 = parse_ast(code2)
+        
+        # Get syntactic features
+        syn1 = get_syntactic_features(ast_tree1)
+        syn2 = get_syntactic_features(ast_tree2)
+        
+        # Get semantic features
+        sem1 = get_semantic_features(code1)
+        sem2 = get_semantic_features(code2)
+        
+        # Calculate hybrid similarity
+        similarities = hybrid_similarity(code1, code2)
+        
+        return {
+            'similarities': similarities,
+            'lexical_features': (lex1, lex2),
+            'syntactic_features': (syn1, syn2),
+            'ast_trees': (ast_tree1, ast_tree2)
+        }
+
+# UI Elements
+st.title("🔍 Advanced Java Code Clone Detector (IJaDataset 2.1)")
+st.markdown("""
+Detect all types of code clones (Type 1-4) using hybrid approach with:
+- **CodeBERT** for semantic analysis
+- **RNN** for sequence modeling
+- **GNN** for AST structural analysis
+""")
+
+# Dataset selector
+selected_pair = None
+if dataset_pairs:
+    pair_options = {f"{i+1}: {pair['type']}": pair for i, pair in enumerate(dataset_pairs)}
+    selected_option = st.selectbox("Select a preloaded example pair:", list(pair_options.keys()))
+    selected_pair = pair_options[selected_option]
+
+# Layout
+col1, col2 = st.columns(2)
+
+with col1:
+    code1 = st.text_area(
+        "First Java Code", 
+        height=300,
+        value=selected_pair["code1"] if selected_pair else "",
+        help="Enter the first Java code snippet"
+    )
+
+with col2:
+    code2 = st.text_area(
+        "Second Java Code", 
+        height=300,
+        value=selected_pair["code2"] if selected_pair else "",
+        help="Enter the second Java code snippet"
+    )
+
+# Threshold sliders
+st.subheader("Detection Thresholds")
+col1, col2, col3 = st.columns(3)
+
+with col1:
+    threshold_type12 = st.slider(
+        "Type 1/2 Threshold",
+        min_value=0.5,
+        max_value=1.0,
+        value=0.9,
+        step=0.01,
+        help="Threshold for exact/syntactic clones"
+    )
+
+with col2:
+    threshold_type3 = st.slider(
+        "Type 3 Threshold",
+        min_value=0.5,
+        max_value=1.0,
+        value=0.8,
+        step=0.01,
+        help="Threshold for near-miss clones"
+    )
+
+with col3:
+    threshold_type4 = st.slider(
+        "Type 4 Threshold",
+        min_value=0.5,
+        max_value=1.0,
+        value=0.7,
+        step=0.01,
+        help="Threshold for semantic clones"
+    )
+
+# Compare button
+if st.button("Compare Code", type="primary"):
+    if tokenizer is None or code_model is None or rnn_model is None or gnn_model is None:
+        st.error("Models failed to load. Please check the logs.")
+    else:
+        result = compare_code(code1, code2)
+        
+        if result is not None:
+            similarities = result['similarities']
+            lex1, lex2 = result['lexical_features']
+            syn1, syn2 = result['syntactic_features']
+            ast_tree1, ast_tree2 = result['ast_trees']
+            
+            # Display results
+            st.subheader("Detection Results")
+            
+            # Determine clone type
+            combined_sim = similarities['combined']
+            clone_type = "No Clone"
+            
+            if combined_sim >= threshold_type12:
+                clone_type = "Type 1/2 Clone (Exact/Near-Exact)"
+            elif combined_sim >= threshold_type3:
+                clone_type = "Type 3 Clone (Near-Miss)"
+            elif combined_sim >= threshold_type4:
+                clone_type = "Type 4 Clone (Semantic)"
+            
+            # Main metrics
+            col1, col2, col3 = st.columns(3)
+            
+            with col1:
+                st.metric("Combined Similarity", f"{combined_sim:.3f}")
+            
+            with col2:
+                st.metric("Detected Clone Type", clone_type)
+            
+            with col3:
+                st.metric("CodeBERT Similarity", f"{similarities['codebert']:.3f}")
+            
+            # Detailed metrics
+            with st.expander("Detailed Similarity Scores"):
+                cols = st.columns(3)
+                with cols[0]:
+                    st.metric("RNN Similarity", f"{similarities['rnn']:.3f}")
+                with cols[1]:
+                    st.metric("GNN Similarity", f"{similarities['gnn']:.3f}")
+                with cols[2]:
+                    st.metric("Lexical Similarity", 
+                            f"{sum(lex1[k] == lex2[k] for k in lex1)/max(len(lex1),1):.2f}")
+            
+            # Feature comparison
+            with st.expander("Feature Analysis"):
+                st.subheader("Lexical Features")
+                lex_df = pd.DataFrame([lex1, lex2], index=["Code 1", "Code 2"])
+                st.dataframe(lex_df)
+                
+                st.subheader("Syntactic Features (AST Node Counts)")
+                syn_df = pd.DataFrame([syn1, syn2], index=["Code 1", "Code 2"]).fillna(0)
+                st.dataframe(syn_df)
+            
+            # AST Visualization
+            if ast_tree1 and ast_tree2:
+                with st.expander("AST Visualization (First 20 nodes)"):
+                    st.write("AST visualization would be implemented here with graphviz")
+                    # In a real implementation, you would use graphviz to render the ASTs
+                    # st.graphviz_chart(ast_to_graphviz(ast_tree1))
+                    # st.graphviz_chart(ast_to_graphviz(ast_tree2))
+            
+            # Normalized code view
+            with st.expander("Show normalized code"):
+                tab1, tab2 = st.tabs(["First Code", "Second Code"])
+                
+                with tab1:
+                    st.code(normalize_code(code1))
+                
+                with tab2:
+                    st.code(normalize_code(code2))
+
+# Footer
+st.markdown("---")
+st.markdown("""
+*Dataset Information*:
+- Using IJaDataset 2.1 from Kaggle
+- Contains 100K Java files with clone annotations
+- Clone types: Type-1, Type-2, Type-3, and Type-4 clones
+
+*Model Architecture*:
+- **CodeBERT**: Pre-trained model for semantic analysis
+- **RNN**: Processes token sequences for sequential patterns
+- **GNN**: Analyzes AST structure for syntactic patterns
+- **Hybrid Approach**: Combines all techniques for comprehensive detection
+""")