Update app.py

app.py

@@ -1,235 +1,255 @@
 import os
-import streamlit as st
+import re
+import time
+import random
+import zipfile
 import javalang
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import re
-import numpy as np
+import torch_geometric
+from torch_geometric.data import Data, Dataset, DataLoader
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import precision_recall_fscore_support
+from tqdm import tqdm
 import networkx as nx
-from transformers import AutoTokenizer, AutoModel
-import warnings
-import pandas as pd
-from collections import defaultdict
-
-# Configuration
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-warnings.filterwarnings("ignore")
-
-# Constants
-MODEL_NAME = "microsoft/codebert-base"
-MAX_LENGTH = 512
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-# Set up page config
-st.set_page_config(
-    page_title="Java Code Clone Detector",
-    page_icon="🔍",
-    layout="wide"
-)
-
-# Simplified RNN Model (for Hugging Face compatibility)
-class SimpleRNN(nn.Module):
-    def __init__(self, input_size=768, hidden_size=128):
-        super().__init__()
-        self.rnn = nn.RNN(input_size, hidden_size, batch_first=True)
-        self.fc = nn.Linear(hidden_size, 1)
-        
-    def forward(self, x):
-        out, _ = self.rnn(x)
-        return torch.sigmoid(self.fc(out[:, -1]))
-
-# Model Loading with caching
-@st.cache_resource(show_spinner=False)
-def load_models():
-    try:
-        with st.spinner('Loading models (first run may take a few minutes)...'):
-            # Load CodeBERT
-            tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-            code_model = AutoModel.from_pretrained(MODEL_NAME).to(DEVICE)
-            
-            # Initialize simple RNN
-            rnn_model = SimpleRNN().to(DEVICE)
-            
-            return tokenizer, code_model, rnn_model
-    except Exception as e:
-        st.error(f"Model loading failed: {str(e)}")
-        return None, None, None
-
-# AST Processing (simplified for Hugging Face)
-def parse_ast(code):
+
+# ---- Utility functions ----
+
+def unzip_dataset(zip_path, extract_to):
+    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+        zip_ref.extractall(extract_to)
+
+def normalize_java_code(code):
+    # Remove single-line comments
+    code = re.sub(r'//.*?\n', '', code)
+    # Remove multi-line comments
+    code = re.sub(r'/\*.*?\*/', '', code, flags=re.DOTALL)
+    # Remove extra spaces and blank lines
+    code = re.sub(r'\s+', ' ', code)
+    return code.strip()
+
+def safe_parse_java(code):
     try:
-        return javalang.parse.parse(code)
-    except:
+        tokens = list(javalang.tokenizer.tokenize(code))
+        parser = javalang.parser.Parser(tokens)
+        tree = parser.parse()
+        return tree
+    except Exception:
         return None
 
-def build_simple_ast_features(ast_tree):
-    if not ast_tree: return {}
-    
-    features = defaultdict(int)
+def ast_to_graph(ast):
+    graph = nx.DiGraph()
     
-    def traverse(node):
-        features[type(node).__name__] += 1
+    def dfs(node, parent_id=None):
+        node_id = len(graph)
+        graph.add_node(node_id, label=type(node).__name__)
+        if parent_id is not None:
+            graph.add_edge(parent_id, node_id)
         for child in getattr(node, 'children', []):
-            if isinstance(child, javalang.ast.Node):
-                traverse(child)
-            elif isinstance(child, (list, tuple)):
+            if isinstance(child, (list, tuple)):
                 for item in child:
                     if isinstance(item, javalang.ast.Node):
-                        traverse(item)
+                        dfs(item, node_id)
+            elif isinstance(child, javalang.ast.Node):
+                dfs(child, node_id)
     
-    traverse(ast_tree)
-    return dict(features)
-
-# Feature Extraction
-def normalize_code(code):
-    code = re.sub(r'//.*?$', '', code, flags=re.MULTILINE)
-    code = re.sub(r'/\*.*?\*/', '', code, flags=re.DOTALL)
-    return re.sub(r'\s+', ' ', code).strip()
+    dfs(ast)
+    return graph
 
-def get_embedding(code, tokenizer, model):
+def tokenize_java_code(code):
     try:
-        inputs = tokenizer(
-            normalize_code(code),
-            return_tensors="pt",
-            truncation=True,
-            max_length=MAX_LENGTH,
-            padding='max_length'
-        ).to(DEVICE)
-        
-        with torch.no_grad():
-            return model(**inputs).last_hidden_state.mean(dim=1)
+        tokens = list(javalang.tokenizer.tokenize(code))
+        token_list = [token.value for token in tokens]
+        return token_list
     except:
-        return None
+        return []
 
-# Similarity Calculations (optimized for Hugging Face)
-def calculate_similarities(code1, code2, models):
-    tokenizer, code_model, rnn_model = models
-    
-    # Get embeddings
-    emb1 = get_embedding(code1, tokenizer, code_model)
-    emb2 = get_embedding(code2, tokenizer, code_model)
-    
-    # Get AST features
-    ast1 = parse_ast(code1)
-    ast2 = parse_ast(code2)
-    ast_features1 = build_simple_ast_features(ast1)
-    ast_features2 = build_simple_ast_features(ast2)
-    
-    # Calculate similarities
-    codebert_sim = 0
-    if emb1 is not None and emb2 is not None:
-        codebert_sim = F.cosine_similarity(emb1, emb2).item()
-    
-    rnn_sim = 0
-    if emb1 is not None and emb2 is not None:
-        with torch.no_grad():
-            rnn_input = torch.cat([emb1, emb2]).unsqueeze(0)
-            rnn_sim = rnn_model(rnn_input).item()
-    
-    # Simple AST similarity (count matching node types)
-    ast_sim = 0
-    if ast_features1 and ast_features2:
-        common_keys = set(ast_features1.keys()) & set(ast_features2.keys())
-        total_keys = set(ast_features1.keys()) | set(ast_features2.keys())
-        ast_sim = len(common_keys) / len(total_keys) if total_keys else 0
-    
-    return {
-        'codebert': codebert_sim,
-        'rnn': rnn_sim,
-        'ast': ast_sim,
-        'combined': 0.5*codebert_sim + 0.3*rnn_sim + 0.2*ast_sim
-    }
-
-# Main UI
-def main():
-    st.title("🔍 Java Code Clone Detector (IJaDataset 2.1)")
-    st.markdown("Detect Type 1-4 clones using hybrid analysis")
-    
-    # Load models
-    models = load_models()
-    if None in models:
-        st.error("Failed to load required models. Please check the logs.")
-        return
-    
-    # Example code pairs
-    example_pairs = {
-        "Type 1 Example": {
-            "code1": "public class Test { public static void main(String[] args) { System.out.println(\"Hello\"); } }",
-            "code2": "public class Test { public static void main(String[] args) { System.out.println(\"Hello\"); } }"
-        },
-        "Type 2 Example": {
-            "code1": "public class Test { public static void main(String[] args) { System.out.println(\"Hello\"); } }",
-            "code2": "public class Example { public static void main(String[] args) { System.out.println(\"Hello\"); } }"
-        },
-        "Type 3 Example": {
-            "code1": "public class Test { public static void main(String[] args) { for(int i=0;i<10;i++) System.out.println(i); } }",
-            "code2": "public class Example { public static void run(String[] params) { for(int j=0;j<10;j++) System.out.println(j); } }"
+# ---- Data Preprocessing ----
+
+class CloneDataset(Dataset):
+    def __init__(self, root_dir, transform=None):
+        super().__init__()
+        self.data_list = []
+        self.labels = []
+        self.skipped_files = 0
+        self.max_tokens = 5000
+
+        clone_dirs = {
+            "Clone_Type1": 1,
+            "Clone_Type2": 1,
+            "Clone_Type3 - ST": 1,
+            "Clone_Type3 - VST": 1,
+            "Clone_Type3 - MT": 0  # Assuming MT = Not Clone
         }
-    }
-    
-    # Code input
-    selected_example = st.selectbox("Select example pair:", list(example_pairs.keys()))
-    
-    col1, col2 = st.columns(2)
-    with col1:
-        code1 = st.text_area(
-            "Code 1", 
-            height=300,
-            value=example_pairs[selected_example]["code1"]
-        )
-    with col2:
-        code2 = st.text_area(
-            "Code 2", 
-            height=300,
-            value=example_pairs[selected_example]["code2"]
-        )
-    
-    # Thresholds
-    st.subheader("Detection Thresholds")
-    cols = st.columns(3)
-    with cols[0]:
-        t1 = st.slider("Type 1/2", 0.85, 1.0, 0.95)
-    with cols[1]:
-        t3 = st.slider("Type 3", 0.7, 0.9, 0.8)
-    with cols[2]:
-        t4 = st.slider("Type 4", 0.5, 0.8, 0.65)
-    
-    # Analysis button
-    if st.button("Analyze Code", type="primary"):
-        with st.spinner("Analyzing code..."):
-            sims = calculate_similarities(code1, code2, models)
-            
-            # Determine clone type
-            clone_type = "No Clone"
-            if sims['combined'] >= t1:
-                clone_type = "Type 1/2 Clone (Exact/Near-Exact)"
-            elif sims['combined'] >= t3:
-                clone_type = "Type 3 Clone (Near-Miss)"
-            elif sims['combined'] >= t4:
-                clone_type = "Type 4 Clone (Semantic)"
-            
-            # Display results
-            st.subheader("Results")
-            
-            # Metrics
-            cols = st.columns(4)
-            cols[0].metric("Combined", f"{sims['combined']:.2f}")
-            cols[1].metric("CodeBERT", f"{sims['codebert']:.2f}")
-            cols[2].metric("RNN", f"{sims['rnn']:.2f}")
-            cols[3].metric("AST", f"{sims['ast']:.2f}")
-            
-            # Progress bar
-            st.progress(sims['combined'])
-            
-            # Final result
-            st.metric("Detection Result", clone_type)
-            
-            # Show details
-            with st.expander("Advanced Details"):
-                st.json(sims)
-                st.code(f"Normalized Code 1:\n{normalize_code(code1)}")
-                st.code(f"Normalized Code 2:\n{normalize_code(code2)}")
+
+        for clone_type, label in clone_dirs.items():
+            clone_path = os.path.join(root_dir, 'Subject_CloneTypes_Directories', clone_type)
+            for root, _, files in os.walk(clone_path):
+                for file in files:
+                    if file.endswith(".java"):
+                        file_path = os.path.join(root, file)
+                        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
+                            code = f.read()
+                            code = normalize_java_code(code)
+                            if len(code.split()) > self.max_tokens:
+                                self.skipped_files += 1
+                                continue
+                            ast = safe_parse_java(code)
+                            if ast is None:
+                                self.skipped_files += 1
+                                continue
+                            graph = ast_to_graph(ast)
+                            tokens = tokenize_java_code(code)
+                            if not tokens:
+                                self.skipped_files += 1
+                                continue
+                            data = {
+                                'graph': graph,
+                                'tokens': tokens,
+                                'label': label
+                            }
+                            self.data_list.append(data)
+
+    def len(self):
+        return len(self.data_list)
+
+    def get(self, idx):
+        data_item = self.data_list[idx]
+        graph = data_item['graph']
+        tokens = data_item['tokens']
+        label = data_item['label']
+
+        # Graph processing
+        edge_index = torch.tensor(list(graph.edges)).t().contiguous()
+
+        node_features = torch.arange(graph.number_of_nodes()).unsqueeze(1).float()
+
+        # Token processing
+        token_indices = torch.tensor([hash(t) % 5000 for t in tokens], dtype=torch.long)
+
+        return edge_index, node_features, token_indices, torch.tensor(label, dtype=torch.long)
+
+# ---- Models ----
+
+class GNNEncoder(nn.Module):
+    def __init__(self, in_channels=1, hidden_dim=64):
+        super().__init__()
+        self.conv1 = torch_geometric.nn.GCNConv(in_channels, hidden_dim)
+        self.conv2 = torch_geometric.nn.GCNConv(hidden_dim, hidden_dim)
+
+    def forward(self, x, edge_index):
+        x = self.conv1(x, edge_index)
+        x = F.relu(x)
+        x = self.conv2(x, edge_index)
+        x = F.relu(x)
+        return torch.mean(x, dim=0)  # Graph-level embedding
+
+class RNNEncoder(nn.Module):
+    def __init__(self, vocab_size=5000, embedding_dim=64, hidden_dim=64):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embedding_dim)
+        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)
+
+    def forward(self, tokens):
+        embeds = self.embedding(tokens)
+        _, (hidden, _) = self.lstm(embeds)
+        return hidden.squeeze(0)
+
+class HybridClassifier(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.gnn = GNNEncoder()
+        self.rnn = RNNEncoder()
+        self.fc = nn.Linear(128, 2)
+
+    def forward(self, edge_index, node_features, tokens):
+        gnn_out = self.gnn(node_features, edge_index)
+        rnn_out = self.rnn(tokens)
+        combined = torch.cat([gnn_out, rnn_out], dim=-1)
+        out = self.fc(combined)
+        return out
+
+# ---- Training and Evaluation ----
+
+def train(model, optimizer, loader, device):
+    model.train()
+    total_loss = 0
+    for edge_index, node_features, tokens, labels in loader:
+        edge_index = edge_index.to(device)
+        node_features = node_features.to(device)
+        tokens = tokens.to(device)
+        labels = labels.to(device)
+
+        optimizer.zero_grad()
+        outputs = model(edge_index, node_features, tokens)
+        loss = F.cross_entropy(outputs.unsqueeze(0), labels.unsqueeze(0))
+        loss.backward()
+        optimizer.step()
+        total_loss += loss.item()
+    return total_loss / len(loader)
+
+def evaluate(model, loader, device):
+    model.eval()
+    preds, labels_all = [], []
+    with torch.no_grad():
+        for edge_index, node_features, tokens, labels in loader:
+            edge_index = edge_index.to(device)
+            node_features = node_features.to(device)
+            tokens = tokens.to(device)
+            labels = labels.to(device)
+
+            outputs = model(edge_index, node_features, tokens)
+            pred = outputs.argmax(dim=-1)
+            preds.append(pred.cpu().numpy())
+            labels_all.append(labels.cpu().numpy())
+
+    preds = np.concatenate(preds)
+    labels_all = np.concatenate(labels_all)
+
+    precision, recall, f1, _ = precision_recall_fscore_support(labels_all, preds, average='binary')
+    return precision, recall, f1
+
+# ---- Main Execution ----
 
 if __name__ == "__main__":
-    main()
+    import numpy as np
+
+    dataset_root = '/content/dataset/archive (1)'
+    unzip_dataset('/content/dataset/archive (1).zip', dataset_root)
+
+    dataset = CloneDataset(dataset_root)
+    print(f"Total valid samples: {dataset.len()}")
+    print(f"Total skipped files: {dataset.skipped_files}")
+
+    indices = list(range(dataset.len()))
+    train_idx, temp_idx = train_test_split(indices, test_size=0.2, random_state=42)
+    val_idx, test_idx = train_test_split(temp_idx, test_size=0.5, random_state=42)
+
+    train_set = torch.utils.data.Subset(dataset, train_idx)
+    val_set = torch.utils.data.Subset(dataset, val_idx)
+    test_set = torch.utils.data.Subset(dataset, test_idx)
+
+    batch_size = 1  # small because of variable graph sizes
+    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
+    val_loader = DataLoader(val_set, batch_size=batch_size)
+    test_loader = DataLoader(test_set, batch_size=batch_size)
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    model = HybridClassifier().to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+    epochs = 5
+
+    start_time = time.time()
+    for epoch in range(epochs):
+        train_loss = train(model, optimizer, train_loader, device)
+        precision, recall, f1 = evaluate(model, val_loader, device)
+        print(f"Epoch {epoch+1}: Loss={train_loss:.4f}, Precision={precision:.4f}, Recall={recall:.4f}, F1={f1:.4f}")
+
+    precision, recall, f1 = evaluate(model, test_loader, device)
+    total_time = time.time() - start_time
+
+    print(f"Test Precision: {precision:.4f}, Recall: {recall:.4f}, F1: {f1:.4f}")
+    print(f"Total execution time: {total_time:.2f} seconds")