Deploy CodeMode via Agent

app.py

@@ -10,114 +10,171 @@ from pathlib import Path
 import chromadb
 from chromadb.config import Settings
 import uuid
+import tempfile
 
-# --- Add scripts to path so we can import ingestion modules ---
-# --- Add scripts to path so we can import ingestion modules ---
-sys.path.append(os.path.dirname(__file__))
+# --- Add scripts to path ---
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
 from scripts.core.ingestion.ingest import GitCrawler
 from scripts.core.ingestion.chunk import RepoChunker
 
 # --- Configuration ---
-MODEL_NAME = "shubharuidas/codebert-base-code-embed-mrl-langchain-langgraph"
+BASELINE_MODEL = "microsoft/codebert-base"
+FINETUNED_MODEL = "shubharuidas/codebert-base-code-embed-mrl-langchain-langgraph"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-DB_DIR = Path("data/chroma_db")
+DB_DIR = Path(os.path.abspath("data/chroma_db_comparison"))
 DB_DIR.mkdir(parents=True, exist_ok=True)
 
-print(f"Loading model: {MODEL_NAME} on {DEVICE}...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModel.from_pretrained(MODEL_NAME)
-model.to(DEVICE)
-model.eval()
-print("Model loaded!")
+print(f"Loading models on {DEVICE}...")
+print("1. Loading baseline model...")
+baseline_tokenizer = AutoTokenizer.from_pretrained(BASELINE_MODEL)
+baseline_model = AutoModel.from_pretrained(BASELINE_MODEL)
+baseline_model.to(DEVICE)
+baseline_model.eval()
 
-# --- Vector Database Setup ---
-# Initialize ChromaDB Client (Persistent)
+print("2. Loading fine-tuned model...")
+finetuned_tokenizer = AutoTokenizer.from_pretrained(FINETUNED_MODEL)
+finetuned_model = AutoModel.from_pretrained(FINETUNED_MODEL)
+finetuned_model.to(DEVICE)
+finetuned_model.eval()
+print("Both models loaded!")
+
+# --- ChromaDB Setup ---
 chroma_client = chromadb.PersistentClient(path=str(DB_DIR))
+baseline_collection = chroma_client.get_or_create_collection(name="baseline_rag", metadata={"hnsw:space": "cosine"})
+finetuned_collection = chroma_client.get_or_create_collection(name="finetuned_rag", metadata={"hnsw:space": "cosine"})
 
-# Create or Get Collection
-# We use cosine similarity space
-collection = chroma_client.get_or_create_collection(name="codemode_rag", metadata={"hnsw:space": "cosine"})
+# --- Embedding Functions ---
+def compute_baseline_embeddings(text_list):
+    if not text_list: return None
+    inputs = baseline_tokenizer(text_list, return_tensors="pt", padding=True, truncation=True, max_length=512).to(DEVICE)
+    with torch.no_grad():
+        out = baseline_model(**inputs)
+        emb = out.last_hidden_state.mean(dim=1)
+        return F.normalize(emb, p=2, dim=1)
 
-# --- Helper Functions ---
-def compute_embeddings(text_list):
-    """Batch compute embeddings"""
+def compute_finetuned_embeddings(text_list):
     if not text_list: return None
-    # Truncate to 512 tokens to avoid errors
-    inputs = tokenizer(text_list, return_tensors="pt", padding=True, truncation=True, max_length=512).to(DEVICE)
+    inputs = finetuned_tokenizer(text_list, return_tensors="pt", padding=True, truncation=True, max_length=512).to(DEVICE)
     with torch.no_grad():
-        out = model(**inputs)
+        out = finetuned_model(**inputs)
         emb = out.last_hidden_state.mean(dim=1)
         return F.normalize(emb, p=2, dim=1)
 
-def reset_db():
-    """Clear database"""
+# --- Reset Functions ---
+def reset_baseline():
+    chroma_client.delete_collection("baseline_rag")
+    global baseline_collection
+    baseline_collection = chroma_client.get_or_create_collection(name="baseline_rag", metadata={"hnsw:space": "cosine"})
+    return "Baseline database reset."
+
+def reset_finetuned():
+    chroma_client.delete_collection("finetuned_rag")
+    global finetuned_collection
+    finetuned_collection = chroma_client.get_or_create_collection(name="finetuned_rag", metadata={"hnsw:space": "cosine"})
+    return "Fine-tuned database reset."
+
+# --- Database Inspector Functions ---
+def list_baseline_files():
+    count = baseline_collection.count()
+    if count == 0:
+        return [["No data indexed yet", "-", "-"]]
+    
     try:
-        chroma_client.delete_collection("codemode_rag")
-        chroma_client.get_or_create_collection(name="codemode_rag", metadata={"hnsw:space": "cosine"})
-        return "Database reset (All embeddings deleted)."
+        data = baseline_collection.get(limit=min(count, 1000), include=["metadatas"])
+        file_stats = {}
+        for meta in data['metadatas']:
+            fname = meta.get("file_name", "unknown")
+            url = meta.get("url", "unknown")
+            if fname not in file_stats:
+                file_stats[fname] = {"count": 0, "url": url}
+            file_stats[fname]["count"] += 1
+        
+        results = [[fname, stats["count"], stats["url"]] for fname, stats in file_stats.items()]
+        return sorted(results, key=lambda x: x[1], reverse=True)
     except Exception as e:
-        return f"Error resetting DB: {e}"
+        return [[f"Error: {str(e)}", "-", "-"]]
 
-def search_codebase(query, top_k=5):
-    """Semantic Search using ChromaDB"""
-    if collection.count() == 0: return []
+def list_finetuned_files():
+    count = finetuned_collection.count()
+    if count == 0:
+        return [["No data indexed yet", "-", "-"]]
     
-    query_emb = compute_embeddings([query])
+    try:
+        data = finetuned_collection.get(limit=min(count, 1000), include=["metadatas"])
+        file_stats = {}
+        for meta in data['metadatas']:
+            fname = meta.get("file_name", "unknown")
+            url = meta.get("url", "unknown")
+            if fname not in file_stats:
+                file_stats[fname] = {"count": 0, "url": url}
+            file_stats[fname]["count"] += 1
+        
+        results = [[fname, stats["count"], stats["url"]] for fname, stats in file_stats.items()]
+        return sorted(results, key=lambda x: x[1], reverse=True)
+    except Exception as e:
+        return [[f"Error: {str(e)}", "-", "-"]]
+
+# --- Search Functions ---
+def search_baseline(query, top_k=5):
+    if baseline_collection.count() == 0: return []
+    query_emb = compute_baseline_embeddings([query])
     if query_emb is None: return []
-    
-    # Convert tensor to list for Chroma
     query_vec = query_emb.cpu().numpy().tolist()[0]
-    
-    results = collection.query(
-        query_embeddings=[query_vec],
-        n_results=min(top_k, collection.count()),
-        include=["metadatas", "documents", "distances"]
-    )
-    
-    # Parse items
+    results = baseline_collection.query(query_embeddings=[query_vec], n_results=min(top_k, baseline_collection.count()), include=["metadatas", "documents", "distances"])
     output = []
     if results['ids']:
         for i in range(len(results['ids'][0])):
             meta = results['metadatas'][0][i]
             code = results['documents'][0][i]
             dist = results['distances'][0][i]
-            score = 1 - dist # Cosine distance to similarity
-            
-            link_icon = "[Link]" if score > 0.7 else ""
-            output.append([meta.get("file_name", "unknown"), f"{score:.4f} {link_icon}", code[:300] + "..."])
-            
+            score = 1 - dist
+            output.append([meta.get("file_name", "unknown"), f"{score:.4f}", code[:300] + "..."])
     return output
 
-def fn_ingest(repo_url):
-    """
-    1. Clone Repo
-    2. Chunk Files
-    3. Compute Embeddings (Batched)
-    4. Store in ChromaDB
-    """
+def search_finetuned(query, top_k=5):
+    if finetuned_collection.count() == 0: return []
+    query_emb = compute_finetuned_embeddings([query])
+    if query_emb is None: return []
+    query_vec = query_emb.cpu().numpy().tolist()[0]
+    results = finetuned_collection.query(query_embeddings=[query_vec], n_results=min(top_k, finetuned_collection.count()), include=["metadatas", "documents", "distances"])
+    output = []
+    if results['ids']:
+        for i in range(len(results['ids'][0])):
+            meta = results['metadatas'][0][i]
+            code = results['documents'][0][i]
+            dist = results['distances'][0][i]
+            score = 1 - dist
+            output.append([meta.get("file_name", "unknown"), f"{score:.4f}", code[:300] + "..."])
+    return output
+
+def search_comparison(query, top_k=5):
+    baseline_results = search_baseline(query, top_k)
+    finetuned_results = search_finetuned(query, top_k)
+    return baseline_results, finetuned_results
+
+# --- Ingestion Functions ---
+def ingest_from_url(repo_url):
     if not repo_url.startswith("http"):
-        return "Invalid URL"
+        yield "Invalid URL"
+        return
     
     DATA_DIR = Path(os.path.abspath("data/raw_ingest"))
     import stat
     def remove_readonly(func, path, _):
         os.chmod(path, stat.S_IWRITE)
         func(path)
-
+    
     try:
-        # Clean up old raw data
         if DATA_DIR.exists():
             shutil.rmtree(DATA_DIR, onerror=remove_readonly)
         
-        # 1. Clone
         yield f"Cloning {repo_url}..."
         crawler = GitCrawler(cache_dir=DATA_DIR)
         repo_path = crawler.clone_repository(repo_url)
-        
         if not repo_path:
-             return "Failed to clone repository."
-
-        # 2. Chunk
+            yield "Failed to clone repository."
+            return
+        
         yield "Listing files..."
         files = crawler.list_files(repo_path, extensions={'.py', '.md', '.json', '.js', '.ts', '.java', '.cpp'})
         if isinstance(files, tuple): files = [f.path for f in files[0]]
@@ -136,65 +193,124 @@ def fn_ingest(repo_url):
                 all_chunks.extend(file_chunks)
             except Exception as e:
                 print(f"Skipping {file_path}: {e}")
-
+        
         if not all_chunks:
-            return "No valid chunks found."
-
-        # 3. Indexing Loop (Batched)
+            yield "No valid chunks found."
+            return
+        
         total_chunks = len(all_chunks)
-        yield f"Generated {total_chunks} chunks. Embedding & Indexing into ChromaDB..."
+        yield f"Generated {total_chunks} chunks. Embedding (BASELINE)..."
         
         batch_size = 64
+        # Index with baseline
         for i in range(0, total_chunks, batch_size):
             batch = all_chunks[i:i+batch_size]
-            
-            # Prepare data
             texts = [c.code for c in batch]
             ids = [str(uuid.uuid4()) for _ in batch]
             metadatas = [{"file_name": Path(c.file_path).name, "url": repo_url} for c in batch]
             
-            # Compute Embeddings
-            embeddings = compute_embeddings(texts)
+            embeddings = compute_baseline_embeddings(texts)
             if embeddings is not None:
-                # Add to Chroma
-                collection.add(
-                    ids=ids,
-                    embeddings=embeddings.cpu().numpy().tolist(),
-                    metadatas=metadatas,
-                    documents=texts
-                )
+                baseline_collection.add(ids=ids, embeddings=embeddings.cpu().numpy().tolist(), metadatas=metadatas, documents=texts)
+            yield f"Baseline: {min(i+batch_size, total_chunks)}/{total_chunks}"
+        
+        yield f"Embedding (FINE-TUNED)..."
+        # Index with fine-tuned
+        for i in range(0, total_chunks, batch_size):
+            batch = all_chunks[i:i+batch_size]
+            texts = [c.code for c in batch]
+            ids = [str(uuid.uuid4()) for _ in batch]
+            metadatas = [{"file_name": Path(c.file_path).name, "url": repo_url} for c in batch]
             
-            progress = int((i / total_chunks) * 100)
-            yield f"Indexed {min(i+batch_size, total_chunks)}/{total_chunks} ({progress}%)"
+            embeddings = compute_finetuned_embeddings(texts)
+            if embeddings is not None:
+                finetuned_collection.add(ids=ids, embeddings=embeddings.cpu().numpy().tolist(), metadatas=metadatas, documents=texts)
+            yield f"Fine-tuned: {min(i+batch_size, total_chunks)}/{total_chunks}"
+        
+        yield f"SUCCESS! Indexed {total_chunks} chunks in both databases."
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        yield f"Error: {str(e)}"
+
+def ingest_from_files(files):
+    if not files or len(files) == 0:
+        yield "No files uploaded."
+        return
+    
+    try:
+        yield f"Processing {len(files)} file(s)..."
+        
+        chunker = RepoChunker()
+        all_chunks = []
+        
+        for i, file in enumerate(files):
+            yield f"Chunking file {i+1}/{len(files)}: {Path(file.name).name}"
+            try:
+                # Gradio file upload: file.name contains the temp path
+                file_path = Path(file.name)
+                meta = {"file_name": file_path.name, "url": "uploaded"}
+                file_chunks = chunker.chunk_file(file_path, repo_metadata=meta)
+                all_chunks.extend(file_chunks)
+            except Exception as e:
+                yield f"Error chunking {Path(file.name).name}: {str(e)}"
+                import traceback
+                traceback.print_exc()
 
-        count = collection.count()
-        yield f"Success! Database now has {count} code chunks. Ready for search."
         
+        if not all_chunks:
+            yield "No valid chunks found."
+            return
+        
+        total_chunks = len(all_chunks)
+        yield f"Generated {total_chunks} chunks. Embedding (BASELINE)..."
+        
+        batch_size = 64
+        for i in range(0, total_chunks, batch_size):
+            batch = all_chunks[i:i+batch_size]
+            texts = [c.code for c in batch]
+            ids = [str(uuid.uuid4()) for _ in batch]
+            metadatas = [{"file_name": Path(c.file_path).name, "url": "uploaded"} for c in batch]
+            
+            embeddings = compute_baseline_embeddings(texts)
+            if embeddings is not None:
+                baseline_collection.add(ids=ids, embeddings=embeddings.cpu().numpy().tolist(), metadatas=metadatas, documents=texts)
+            yield f"Baseline: {min(i+batch_size, total_chunks)}/{total_chunks}"
+        
+        yield f"Embedding (FINE-TUNED)..."
+        for i in range(0, total_chunks, batch_size):
+            batch = all_chunks[i:i+batch_size]
+            texts = [c.code for c in batch]
+            ids = [str(uuid.uuid4()) for _ in batch]
+            metadatas = [{"file_name": Path(c.file_path).name, "url": "uploaded"} for c in batch]
+            
+            embeddings = compute_finetuned_embeddings(texts)
+            if embeddings is not None:
+                finetuned_collection.add(ids=ids, embeddings=embeddings.cpu().numpy().tolist(), metadatas=metadatas, documents=texts)
+            yield f"Fine-tuned: {min(i+batch_size, total_chunks)}/{total_chunks}"
+        
+        yield f"SUCCESS! Indexed {total_chunks} chunks from uploaded files."
     except Exception as e:
         import traceback
         traceback.print_exc()
         yield f"Error: {str(e)}"
 
-# --- Analysis Functions ---
-def fn_analyze_embeddings():
-    count = collection.count()
+# --- Analysis & Evaluation Functions ---
+def analyze_embeddings_baseline():
+    count = baseline_collection.count()
     if count < 5:
         return "Not enough data (Need > 5 chunks).", None
     
     try:
-        # Fetch all embeddings (Limit to 2000 for visualization speed)
         limit = min(count, 2000)
-        data = collection.get(limit=limit, include=["embeddings", "metadatas"])
+        data = baseline_collection.get(limit=limit, include=["embeddings", "metadatas"])
         
         X = torch.tensor(data['embeddings'])
-        
-        # PCA
         X_mean = torch.mean(X, 0)
         X_centered = X - X_mean
         U, S, V = torch.pca_lowrank(X_centered, q=2)
         projected = torch.matmul(X_centered, V[:, :2]).numpy()
         
-        # Diversity
         indices = torch.randint(0, len(X), (min(100, len(X)),))
         sample = X[indices]
         sim_matrix = torch.mm(sample, sample.t())
@@ -203,10 +319,11 @@ def fn_analyze_embeddings():
         diversity_score = 1.0 - avg_sim
         
         metrics = (
+            f"BASELINE MODEL\n"
             f"Total Chunks: {count}\n"
-            f"Analyzed: {len(X)} (Sampled)\n"
+            f"Analyzed: {len(X)}\n"
             f"Diversity Score: {diversity_score:.4f}\n"
-            f"Est. Avg Similarity: {avg_sim:.4f}"
+            f"Avg Similarity: {avg_sim:.4f}"
         )
         
         plot_df = pd.DataFrame({
@@ -215,22 +332,61 @@ def fn_analyze_embeddings():
             "topic": [m.get("file_name", "unknown") for m in data['metadatas']]
         })
         
-        return metrics, gr.ScatterPlot(value=plot_df, x="x", y="y", color="topic", title="Semantic Space", tooltip="topic")
+        return metrics, gr.ScatterPlot(value=plot_df, x="x", y="y", color="topic", title="Baseline Semantic Space", tooltip="topic")
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        return f"Error: {e}", None
+
+def analyze_embeddings_finetuned():
+    count = finetuned_collection.count()
+    if count < 5:
+        return "Not enough data (Need > 5 chunks).", None
+    
+    try:
+        limit = min(count, 2000)
+        data = finetuned_collection.get(limit=limit, include=["embeddings", "metadatas"])
+        
+        X = torch.tensor(data['embeddings'])
+        X_mean = torch.mean(X, 0)
+        X_centered = X - X_mean
+        U, S, V = torch.pca_lowrank(X_centered, q=2)
+        projected = torch.matmul(X_centered, V[:, :2]).numpy()
+        
+        indices = torch.randint(0, len(X), (min(100, len(X)),))
+        sample = X[indices]
+        sim_matrix = torch.mm(sample, sample.t())
+        mask = ~torch.eye(len(sample), dtype=bool)
+        avg_sim = sim_matrix[mask].mean().item()
+        diversity_score = 1.0 - avg_sim
+        
+        metrics = (
+            f"FINE-TUNED MODEL\n"
+            f"Total Chunks: {count}\n"
+            f"Analyzed: {len(X)}\n"
+            f"Diversity Score: {diversity_score:.4f}\n"
+            f"Avg Similarity: {avg_sim:.4f}"
+        )
+        
+        plot_df = pd.DataFrame({
+            "x": projected[:, 0],
+            "y": projected[:, 1],
+            "topic": [m.get("file_name", "unknown") for m in data['metadatas']]
+        })
         
+        return metrics, gr.ScatterPlot(value=plot_df, x="x", y="y", color="topic", title="Fine-tuned Semantic Space", tooltip="topic")
     except Exception as e:
         import traceback
         traceback.print_exc()
-        return f"Analysis Error: {e}", None
+        return f"Error: {e}", None
 
-def fn_evaluate_retrieval(sample_limit):
-    count = collection.count()
+def evaluate_retrieval_baseline(sample_limit):
+    count = baseline_collection.count()
     if count < 10: return "Not enough data for evaluation (Need > 10 chunks)."
     
     try:
-        # Sample random chunks
-        # Chroma doesn't support random sample easily, so we get a larger batch and pick random
-        fetch_limit = min(count, 2000) # Fetch up to 2k to sample from
-        data = collection.get(limit=fetch_limit, include=["documents"])
+        fetch_limit = min(count, 2000)
+        data = baseline_collection.get(limit=fetch_limit, include=["documents"])
         
         import random
         actual_sample_size = min(sample_limit, len(data['ids']))
@@ -240,191 +396,210 @@ def fn_evaluate_retrieval(sample_limit):
         hits_at_5 = 0
         mrr_sum = 0
         
-        # Generator for progress updates
-        yield f"Running evaluation on {actual_sample_size} chunks..."
+        yield f"BASELINE: Evaluating {actual_sample_size} chunks..."
         
         for i, idx in enumerate(sample_indices):
             target_id = data['ids'][idx]
             code = data['documents'][idx]
-            
-            # Synthetic Query
             query = "\n".join(code.split("\n")[:3])
-            query_emb = compute_embeddings([query]).cpu().numpy().tolist()[0]
-            
-            # Query DB
-            results = collection.query(query_embeddings=[query_emb], n_results=10)
-            
-            # Check results
+            query_emb = compute_baseline_embeddings([query]).cpu().numpy().tolist()[0]
+            results = baseline_collection.query(query_embeddings=[query_emb], n_results=10)
             found_ids = results['ids'][0]
             if target_id in found_ids:
                 rank = found_ids.index(target_id) + 1
                 mrr_sum += 1.0 / rank
                 if rank == 1: hits_at_1 += 1
                 if rank <= 5: hits_at_5 += 1
-            
             if i % 10 == 0:
-                yield f"Evaluated {i}/{actual_sample_size}..."
+                yield f"Baseline: {i}/{actual_sample_size}..."
         
         recall_1 = hits_at_1 / actual_sample_size
         recall_5 = hits_at_5 / actual_sample_size
         mrr = mrr_sum / actual_sample_size
         
         report = (
-            f"Evaluation on {actual_sample_size} random chunks:\n"
-            f"--------------------------------------------\n"
+            f"BASELINE EVALUATION ({actual_sample_size} chunks)\n"
+            f"{'='*40}\n"
             f"Recall@1: {recall_1:.4f}\n"
             f"Recall@5: {recall_5:.4f}\n"
-            f"MRR:      {mrr:.4f}\n"
-            f"\n(Note: Using ChromaDB for retrieval)"
+            f"MRR:      {mrr:.4f}"
         )
         yield report
     except Exception as e:
         import traceback
         traceback.print_exc()
-        yield f"Eval Error: {e}"
+        yield f"Error: {e}"
 
+def evaluate_retrieval_finetuned(sample_limit):
+    count = finetuned_collection.count()
+    if count < 10: return "Not enough data for evaluation (Need > 10 chunks)."
+    
+    try:
+        fetch_limit = min(count, 2000)
+        data = finetuned_collection.get(limit=fetch_limit, include=["documents"])
+        
+        import random
+        actual_sample_size = min(sample_limit, len(data['ids']))
+        sample_indices = random.sample(range(len(data['ids'])), actual_sample_size)
+        
+        hits_at_1 = 0
+        hits_at_5 = 0
+        mrr_sum = 0
+        
+        yield f"FINE-TUNED: Evaluating {actual_sample_size} chunks..."
+        
+        for i, idx in enumerate(sample_indices):
+            target_id = data['ids'][idx]
+            code = data['documents'][idx]
+            query = "\n".join(code.split("\n")[:3])
+            query_emb = compute_finetuned_embeddings([query]).cpu().numpy().tolist()[0]
+            results = finetuned_collection.query(query_embeddings=[query_emb], n_results=10)
+            found_ids = results['ids'][0]
+            if target_id in found_ids:
+                rank = found_ids.index(target_id) + 1
+                mrr_sum += 1.0 / rank
+                if rank == 1: hits_at_1 += 1
+                if rank <= 5: hits_at_5 += 1
+            if i % 10 == 0:
+                yield f"Fine-tuned: {i}/{actual_sample_size}..."
+        
+        recall_1 = hits_at_1 / actual_sample_size
+        recall_5 = hits_at_5 / actual_sample_size
+        mrr = mrr_sum / actual_sample_size
+        
+        report = (
+            f"FINE-TUNED EVALUATION ({actual_sample_size} chunks)\n"
+            f"{'='*40}\n"
+            f"Recall@1: {recall_1:.4f}\n"
+            f"Recall@5: {recall_5:.4f}\n"
+            f"MRR:      {mrr:.4f}"
+        )
+        yield report
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        yield f"Error: {e}"
 
-# --- UI Layout ---
-theme = gr.themes.Soft(
-    primary_hue="slate",
-    neutral_hue="slate",
-    spacing_size="sm",
-    radius_size="md"
-).set(
-    body_background_fill="*neutral_50",
-    block_background_fill="white",
-    block_border_width="1px",
-    block_title_text_weight="600"
-)
+# --- UI ---
+theme = gr.themes.Soft(primary_hue="slate", neutral_hue="slate", spacing_size="sm", radius_size="md").set(body_background_fill="*neutral_50", block_background_fill="white", block_border_width="1px", block_title_text_weight="600")
 
 css = """
-h1 {
-    text-align: center;
-    font-family: 'Inter', sans-serif;
-    margin-bottom: 1rem;
-    color: #1e293b;
-}
-.gradio-container {
-    max-width: 1200px !important;
-    margin: auto;
-}
+h1 { text-align: center; font-family: 'Inter', sans-serif; margin-bottom: 1rem; color: #1e293b; }
+.gradio-container { max-width: 1400px !important; margin: auto; }
+.comparison-header { font-size: 1.1rem; font-weight: 600; color: #334155; text-align: center; padding: 0.5rem; }
 """
 
-with gr.Blocks(theme=theme, css=css, title="CodeMode") as demo:
-    gr.Markdown("# CodeMode")
+with gr.Blocks(theme=theme, css=css, title="CodeMode - Baseline vs Fine-tuned") as demo:
+    gr.Markdown("# CodeMode: Baseline vs Fine-tuned Model Comparison")
+    gr.Markdown("Compare retrieval performance between **microsoft/codebert-base** (baseline) and **MRL-enhanced fine-tuned** model")
     
     with gr.Tabs():
-        # --- TAB 1: INGEST ---
-        with gr.Tab("1. Ingest GitHub Repo"):
-            gr.Markdown("### Connect a Repository")
-            with gr.Row():
-                repo_input = gr.Textbox(label="GitHub URL", placeholder="https://github.com/fastapi/fastapi", value="https://github.com/langchain-ai/langgraph")
-                ingest_btn = gr.Button("Ingest & Index", variant="primary")
+        # TAB 1: INGEST
+        with gr.Tab("1. Ingest Code"):
+            with gr.Tabs():
+                with gr.Tab("GitHub Repository"):
+                    repo_input = gr.Textbox(label="GitHub URL", placeholder="https://github.com/pallets/flask")
+                    ingest_url_btn = gr.Button("Ingest from URL", variant="primary")
+                    url_status = gr.Textbox(label="Status")
+                    ingest_url_btn.click(ingest_from_url, inputs=repo_input, outputs=url_status)
+                
+                with gr.Tab("Upload Python Files"):
+                    file_upload = gr.File(label="Upload .py files", file_types=[".py"], file_count="multiple")
+                    ingest_files_btn = gr.Button("Ingest Uploaded Files", variant="primary")
+                    upload_status = gr.Textbox(label="Status")
+                    ingest_files_btn.click(ingest_from_files, inputs=file_upload, outputs=upload_status)
             
             with gr.Row():
-                 reset_btn = gr.Button("Reset Database", variant="stop")
-                 ingest_status = gr.Textbox(label="Status")
+                reset_baseline_btn = gr.Button("Reset Baseline DB", variant="stop")
+                reset_finetuned_btn = gr.Button("Reset Fine-tuned DB", variant="stop")
+                reset_status = gr.Textbox(label="Reset Status")
+            
+            reset_baseline_btn.click(reset_baseline, inputs=[], outputs=reset_status)
+            reset_finetuned_btn.click(reset_finetuned, inputs=[], outputs=reset_status)
+            
+            gr.Markdown("---")
+            gr.Markdown("### Database Inspector")
+            gr.Markdown("View indexed files in each collection")
             
-            with gr.Accordion("Database Inspector", open=False):
-                list_files_btn = gr.Button("Refresh File List")
-                files_df = gr.Dataframe(
-                    headers=["File Name", "Chunks", "Source URL"],
-                    datatype=["str", "number", "str"],
-                    interactive=False
-                )
-
-            def fn_list_files():
-                count = collection.count()
-                if count == 0: return [["Database Empty", 0, "-"]]
-                
-                try:
-                    # Fetch all metadata (limit to 10k to prevent UI freeze)
-                    limit = min(count, 10000)
-                    data = collection.get(limit=limit, include=["metadatas"])
-                    
-                    if not data or 'metadatas' not in data or data['metadatas'] is None:
-                         return [["Error: No metadata found", 0, "-"]]
-
-                    # Aggregate stats
-                    file_counts = {} # filename -> count
-                    file_urls = {} # filename -> url
-                    
-                    for meta in data['metadatas']:
-                        if meta is None: continue # Skip None entries
-                        fname = meta.get("file_name", "unknown")
-                        url = meta.get("url", "-")
-                        file_counts[fname] = file_counts.get(fname, 0) + 1
-                        file_urls[fname] = url
-                    
-                    # Convert to list
-                    output = []
-                    for fname, count in file_counts.items():
-                        output.append([fname, count, file_urls[fname]])
-                    
-                    if not output:
-                        return [["No files found in metadata", 0, "-"]]
-                        
-                    # Sort by chunk count (descending)
-                    output.sort(key=lambda x: x[1], reverse=True)
-                    return output
-                except Exception as e:
-                    import traceback
-                    traceback.print_exc()
-                    return [[f"Error: {str(e)}", 0, "-"]]
-
-            ingest_btn.click(fn_ingest, inputs=repo_input, outputs=[ingest_status])
-            reset_btn.click(fn=reset_db, inputs=[], outputs=[ingest_status])
-            list_files_btn.click(fn_list_files, inputs=[], outputs=[files_df])
-
-        # --- TAB 2: SEARCH ---
-        with gr.Tab("2. Semantic Search"):
-            gr.Markdown("### Search the Ingested Code")
             with gr.Row():
-                search_box = gr.Textbox(label="Search Query", placeholder="e.g., 'how to create a state graph'")
-                search_btn = gr.Button("Search", variant="primary")
+                with gr.Column():
+                    gr.Markdown("#### Baseline Collection")
+                    inspect_baseline_btn = gr.Button("Inspect Baseline DB", variant="secondary")
+                    baseline_files_df = gr.Dataframe(
+                        headers=["File Name", "Chunks", "Source URL"],
+                        datatype=["str", "number", "str"],
+                        interactive=False,
+                        value=[["No data yet", "-", "-"]]
+                    )
+                    inspect_baseline_btn.click(list_baseline_files, inputs=[], outputs=baseline_files_df)
+                
+                with gr.Column():
+                    gr.Markdown("#### Fine-tuned Collection")
+                    inspect_finetuned_btn = gr.Button("Inspect Fine-tuned DB", variant="secondary")
+                    finetuned_files_df = gr.Dataframe(
+                        headers=["File Name", "Chunks", "Source URL"],
+                        datatype=["str", "number", "str"],
+                        interactive=False,
+                        value=[["No data yet", "-", "-"]]
+                    )
+                    inspect_finetuned_btn.click(list_finetuned_files, inputs=[], outputs=finetuned_files_df)
+        
+        # TAB 2: COMPARISON SEARCH
+        with gr.Tab("2. Comparison Search (Note: Semantic search is sensitive to query phrasing)"):
+            gr.Markdown("### Side-by-Side Retrieval Comparison")
+            search_query = gr.Textbox(label="Search Query", placeholder="e.g., 'Flask route decorator'")
+            compare_btn = gr.Button("Compare Models", variant="primary")
             
-            results_df = gr.Dataframe(
-                headers=["File Name", "Score", "Code Snippet"],
-                datatype=["str", "str", "str"],
-                interactive=False,
-                wrap=True
-            )
-            search_btn.click(fn=search_codebase, inputs=search_box, outputs=results_df)
-
-        # --- TAB 3: CODE SEARCH ---
-        with gr.Tab("3. Find Similar Code"):
-            gr.Markdown("### Code-to-Code Retrieval")
             with gr.Row():
-                code_input = gr.Code(label="Reference Code", language="python")
-                code_search_btn = gr.Button("Find Matches", variant="primary")
+                with gr.Column():
+                    gr.Markdown("<div class='comparison-header'>BASELINE (CodeBERT)</div>", elem_classes="comparison-header")
+                    baseline_results = gr.Dataframe(headers=["File", "Score", "Code Snippet"], datatype=["str", "str", "str"], interactive=False, wrap=True)
                 
-            code_results_df = gr.Dataframe(
-                headers=["File Name", "Score", "Matched Code"],
-                datatype=["str", "str", "str"],
-                interactive=False,
-                wrap=True
-            )
-            code_search_btn.click(fn=search_codebase, inputs=code_input, outputs=code_results_df)
-
-        # --- TAB 4: MLOps MONITORING ---
-        with gr.Tab("4. Deployment Monitoring"):
+                with gr.Column():
+                    gr.Markdown("<div class='comparison-header'>FINE-TUNED (MRL-Enhanced)</div>", elem_classes="comparison-header")
+                    finetuned_results = gr.Dataframe(headers=["File", "Score", "Code Snippet"], datatype=["str", "str", "str"], interactive=False, wrap=True)
+            
+            compare_btn.click(search_comparison, inputs=search_query, outputs=[baseline_results, finetuned_results])
+        
+        
+        # TAB 3: DEPLOYMENT MONITORING
+        with gr.Tab("3. Deployment Monitoring"):
             gr.Markdown("### Embedding Quality Analysis")
-            analyze_btn = gr.Button("Analyze Embeddings", variant="secondary")
+            gr.Markdown("Analyze the semantic space distribution and diversity of embeddings")
             
             with gr.Row():
-                quality_metrics = gr.Textbox(label="Quality Metrics")
-                plot_output = gr.ScatterPlot(label="Semantic Space (PCA)")
+                with gr.Column():
+                    gr.Markdown("#### Baseline Model")
+                    analyze_baseline_btn = gr.Button("Analyze Baseline Embeddings", variant="secondary")
+                    baseline_metrics = gr.Textbox(label="Baseline Metrics")
+                    baseline_plot = gr.ScatterPlot(label="Baseline Semantic Space (PCA)")
+                    analyze_baseline_btn.click(analyze_embeddings_baseline, inputs=[], outputs=[baseline_metrics, baseline_plot])
+                
+                with gr.Column():
+                    gr.Markdown("#### Fine-tuned Model")
+                    analyze_finetuned_btn = gr.Button("Analyze Fine-tuned Embeddings", variant="secondary")
+                    finetuned_metrics = gr.Textbox(label="Fine-tuned Metrics")
+                    finetuned_plot = gr.ScatterPlot(label="Fine-tuned Semantic Space (PCA)")
+                    analyze_finetuned_btn.click(analyze_embeddings_finetuned, inputs=[], outputs=[finetuned_metrics, finetuned_plot])
             
-            analyze_btn.click(fn_analyze_embeddings, inputs=[], outputs=[quality_metrics, plot_output])
-
-            gr.Markdown("### Extrinsic Evaluation (Retrieval Performance)")
-            with gr.Row():
-                eval_size = gr.Slider(minimum=10, maximum=1000, value=50, step=10, label="Sample Size (Chunks)")
-                eval_btn = gr.Button("Run Retrieval Evaluation", variant="primary")
+            gr.Markdown("---")
+            gr.Markdown("### Retrieval Performance Evaluation")
+            gr.Markdown("Evaluate retrieval accuracy using synthetic queries (query = first 3 lines of code)")
             
-            eval_output = gr.Textbox(label="Evaluation Report")
+            eval_size = gr.Slider(minimum=10, maximum=500, value=50, step=10, label="Sample Size (Chunks to Evaluate)")
             
-            eval_btn.click(fn_evaluate_retrieval, inputs=[eval_size], outputs=eval_output)
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("#### Baseline Evaluation")
+                    eval_baseline_btn = gr.Button("Run Baseline Evaluation", variant="primary")
+                    baseline_eval_output = gr.Textbox(label="Baseline Results")
+                    eval_baseline_btn.click(evaluate_retrieval_baseline, inputs=[eval_size], outputs=baseline_eval_output)
+                
+                with gr.Column():
+                    gr.Markdown("#### Fine-tuned Evaluation")
+                    eval_finetuned_btn = gr.Button("Run Fine-tuned Evaluation", variant="primary")
+                    finetuned_eval_output = gr.Textbox(label="Fine-tuned Results")
+                    eval_finetuned_btn.click(evaluate_retrieval_finetuned, inputs=[eval_size], outputs=finetuned_eval_output)
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=False)