Update app with better diff, new style

app/detector.py

@@ -3,7 +3,7 @@ import json
 import math
 import os
 import re
-from typing import Callable
+from typing import Callable, Iterable
 from functools import cache
 from pathlib import Path
 
@@ -15,7 +15,18 @@ from transformers import AutoModel, AutoTokenizer
 
 import transformers
 
-MODELS_ROOT = Path(transformers.__file__).resolve().parent / "models"
+_LIB_PATH = Path(transformers.__file__).resolve().parent
+_ENV_REPO = os.getenv("TRANSFORMERS_REPO")
+if _ENV_REPO:
+    _env_path = Path(_ENV_REPO)
+    _candidate = _env_path / "src" / "transformers" / "models"
+    if _candidate.exists():
+        MODELS_ROOT = _candidate
+    else:
+        _fallback = _env_path / "models"
+        MODELS_ROOT = _fallback if _fallback.exists() else _LIB_PATH / "models"
+else:
+    MODELS_ROOT = _LIB_PATH / "models"
 
 EMBEDDING_MODEL = "Qwen/Qwen3-Embedding-0.6B"
 BATCH_SIZE = 16
@@ -24,28 +35,46 @@ HUB_DATASET_DEFAULT = "Molbap/modular-detector-embeddings"
 BOILERPLATE_NAMES = {
     "__init__",
     "_init_weights",
+    "__repr__",
+    "extra_repr",
     "get_input_embeddings",
     "set_input_embeddings",
+    "get_output_embeddings",
+    "set_output_embeddings",
+    "tie_weights",
+    "post_init",
+    "forward",
+    "init_weights",
+    "reset_parameters",
+    "training",
 }
 
+_RE_COMMENT = re.compile(r"#.*")
+_RE_IMPORT = re.compile(r"\s*(from|import)\s+")
+_RE_MODEL_HINT = re.compile(r"\d+")
+_RE_LEADING_PREFIX = re.compile(r"^([A-Z][a-z0-9]+)")
+_RE_ALPHANUM = re.compile(r"^([A-Za-z0-9]+)")
+_RE_NORMALIZE = re.compile(r"[^a-z0-9]+")
+
 
 def _sanitize_for_embedding(code: str, model_hint: str | None, symbol_hint: str | None) -> str:
     code = _strip_docstrings(code)
+    cleaned = _RE_COMMENT.sub("", code)
     base = "\n".join(
-        line for line in re.sub(r"#.*", "", code).splitlines() if not re.match(r"\s*(from|import)\s+", line)
+        line for line in cleaned.splitlines() if line.strip() and not _RE_IMPORT.match(line)
     )
     variants = set()
     if model_hint:
         variants.add(model_hint)
         variants.add(model_hint.replace("_", ""))
-        variants.add(re.sub(r"\d+", "", model_hint))
+        variants.add(_RE_MODEL_HINT.sub("", model_hint))
     if symbol_hint:
-        match = re.match(r"^([A-Z][a-z0-9]+)", symbol_hint) or re.match(r"^([A-Za-z0-9]+)", symbol_hint)
+        match = _RE_LEADING_PREFIX.match(symbol_hint) or _RE_ALPHANUM.match(symbol_hint)
         prefix = match.group(1) if match else ""
         if prefix:
             variants.add(prefix)
             variants.add(prefix.replace("_", ""))
-            variants.add(re.sub(r"\d+", "", prefix))
+            variants.add(_RE_MODEL_HINT.sub("", prefix))
     variants |= {variant.lower() for variant in list(variants)}
     sanitized = base
     for variant in sorted({x for x in variants if len(x) >= 3}, key=len, reverse=True):
@@ -54,11 +83,11 @@ def _sanitize_for_embedding(code: str, model_hint: str | None, symbol_hint: str
 
 
 def _normalize(value: str | None) -> str:
-    return re.sub(r"[^a-z0-9]+", "", value.lower()) if value else ""
+    return _RE_NORMALIZE.sub("", value.lower()) if value else ""
 
 
 def _leading_prefix(name: str) -> str:
-    match = re.match(r"^([A-Z][a-z0-9]+)", name) or re.match(r"^([A-Za-z0-9]+)", name)
+    match = _RE_LEADING_PREFIX.match(name) or _RE_ALPHANUM.match(name)
     return match.group(1) if match else ""
 
 
@@ -95,11 +124,22 @@ def _infer_model_prefixes(definitions_kind: dict[str, str]) -> set[str]:
     return prefixes
 
 
-def _calculate_reconstruction_score(scores: list[float], coverage_ratio: float) -> float:
-    if not scores:
-        return 0.0
-    avg_sim = float(sum(scores)) / float(len(scores))
-    return avg_sim * (coverage_ratio**0.5)
+def _calculate_reconstruction_score(
+    contributors: Iterable[dict[str, object]],
+    query_method_count: int,
+) -> tuple[float, int]:
+    if query_method_count <= 0:
+        return 0.0, 0
+    best_scores: dict[str, float] = {}
+    for contributor in contributors:
+        query_name = contributor.get("query")
+        if not query_name:
+            continue
+        score = float(contributor.get("score", 0.0))
+        if score > best_scores.get(str(query_name), 0.0):
+            best_scores[str(query_name)] = score
+    total_similarity = sum(best_scores.values())
+    return total_similarity / float(query_method_count), len(best_scores)
 
 
 def _normalize_source_path(path: str | None) -> str | None:
@@ -144,6 +184,43 @@ def _strip_docstrings(source: str) -> str:
         return source
 
 
+def _strip_docstrings_in_tree(tree: ast.AST) -> None:
+    def strip_in_body(body: list[ast.stmt]) -> None:
+        if not body:
+            return
+        first = body[0]
+        if isinstance(first, ast.Expr) and isinstance(getattr(first, "value", None), ast.Constant):
+            if isinstance(first.value.value, str):
+                body.pop(0)
+
+    strip_in_body(tree.body)
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+            strip_in_body(node.body)
+
+
+def _sanitize_unparsed_code(code: str, model_hint: str | None, symbol_hint: str | None) -> str:
+    cleaned = _RE_COMMENT.sub("", code)
+    base = "\n".join(line for line in cleaned.splitlines() if line.strip() and not _RE_IMPORT.match(line))
+    hints: set[str] = set()
+    if model_hint:
+        hints.add(model_hint)
+        hints.add(model_hint.replace("_", ""))
+        hints.add(_RE_MODEL_HINT.sub("", model_hint))
+    if symbol_hint:
+        match = _RE_LEADING_PREFIX.match(symbol_hint) or _RE_ALPHANUM.match(symbol_hint)
+        prefix = match.group(1) if match else ""
+        if prefix:
+            hints.add(prefix)
+            hints.add(prefix.replace("_", ""))
+            hints.add(_RE_MODEL_HINT.sub("", prefix))
+    hints = {h for h in hints if len(h) >= 3}
+    if hints:
+        pattern = re.compile("|".join(re.escape(h) for h in sorted(hints, key=len, reverse=True)), re.IGNORECASE)
+        base = pattern.sub("Model", base)
+    return base
+
+
 def _normalize_code_for_compare(source: str) -> str:
     stripped = _strip_docstrings(source)
     return "".join(line.strip() for line in stripped.splitlines() if line.strip())
@@ -449,6 +526,8 @@ class CodeSimilarityAnalyzer:
         definitions_kind: dict[str, str] = {}
         lines = code.splitlines()
         tree = ast.parse(code)
+        entries: list[tuple[str, str, ast.AST, ast.ClassDef | None]] = []
+
         for node in ast.iter_child_nodes(tree):
             if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and granularity in ("definition", "method"):
                 segment = ast.get_source_segment(code, node)
@@ -460,9 +539,8 @@ class CodeSimilarityAnalyzer:
                     continue
                 identifier = node.name
                 definitions_raw[identifier] = segment
-                sanitized = _sanitize_for_embedding(segment, model_hint, node.name)
-                definitions_sanitized[identifier] = sanitized
                 definitions_kind[identifier] = "function"
+                entries.append((identifier, "function", node, None))
                 continue
 
             if isinstance(node, ast.ClassDef):
@@ -471,19 +549,14 @@ class CodeSimilarityAnalyzer:
                     start = max(0, node.lineno - 1)
                     end = node.end_lineno
                     class_segment = "\n".join(lines[start:end])
-                class_header = ""
-                if class_segment:
-                    class_header = class_segment.splitlines()[0].strip()
-                class_context = class_header
+                if not class_segment:
+                    continue
 
                 if granularity == "definition":
-                    if not class_segment:
-                        continue
                     identifier = node.name
                     definitions_raw[identifier] = class_segment
-                    sanitized = _sanitize_for_embedding(class_segment, model_hint, node.name)
-                    definitions_sanitized[identifier] = sanitized
                     definitions_kind[identifier] = "class"
+                    entries.append((identifier, "class", node, None))
                     continue
 
                 for child in node.body:
@@ -497,12 +570,24 @@ class CodeSimilarityAnalyzer:
                     if not segment:
                         continue
                     method_name = child.name
-                    combined = f"{class_context}\n{segment}" if class_context else segment
                     identifier = f"{node.name}.{method_name}"
                     definitions_raw[identifier] = segment
-                    sanitized = _sanitize_for_embedding(combined, model_hint, node.name)
-                    definitions_sanitized[identifier] = sanitized
                     definitions_kind[identifier] = "method"
+                    entries.append((identifier, "method", child, node))
+
+        _strip_docstrings_in_tree(tree)
+        for identifier, kind, node, parent in entries:
+            try:
+                if kind == "method" and parent is not None:
+                    parent_header = ast.unparse(parent).splitlines()[0]
+                    combined = f"{parent_header}\n{ast.unparse(node)}"
+                    sanitized = _sanitize_unparsed_code(combined, model_hint, parent.name)
+                else:
+                    sanitized = _sanitize_unparsed_code(ast.unparse(node), model_hint, identifier)
+            except Exception:
+                sanitized = definitions_raw.get(identifier, "")
+            definitions_sanitized[identifier] = sanitized
+
         return definitions_raw, definitions_sanitized, definitions_kind
 
     def analyze_code(
@@ -578,6 +663,15 @@ class CodeSimilarityAnalyzer:
             }
             for identifier, score in candidates:
                 relative_path, match_name = identifier.split(":", 1)
+                is_identical = False
+                match_segment = None
+                if exclude_identical or query_compare:
+                    match_segment = _get_definition_segment(relative_path, match_name)
+                if match_segment is not None:
+                    match_norm = _normalize_code_for_compare(match_segment)
+                    query_norm = query_compare.get(query_identifier)
+                    if query_norm and match_norm == query_norm:
+                        is_identical = True
                 if len(entry_all["embedding"]) < top_k_per_item:
                     full_path, line = _resolve_definition_location(relative_path, match_name)
                     entry_all["embedding"].append(
@@ -588,16 +682,12 @@ class CodeSimilarityAnalyzer:
                             "score": score,
                             "full_path": full_path,
                             "line": line,
+                            "is_identical": is_identical,
                         }
                     )
-                if exclude_identical:
-                    match_segment = _get_definition_segment(relative_path, match_name)
-                    if match_segment is not None:
-                        match_norm = _normalize_code_for_compare(match_segment)
-                        query_norm = query_compare.get(query_identifier)
-                        if query_norm and match_norm == query_norm:
-                            identical_filtered += 1
-                            continue
+                if exclude_identical and is_identical:
+                    identical_filtered += 1
+                    continue
                 full_path, line = _resolve_definition_location(relative_path, match_name)
                 entry["embedding"].append(
                     {
@@ -607,6 +697,7 @@ class CodeSimilarityAnalyzer:
                         "score": score,
                         "full_path": full_path,
                         "line": line,
+                        "is_identical": is_identical,
                     }
                 )
                 if len(entry["embedding"]) >= top_k_per_item and len(entry_all["embedding"]) >= top_k_per_item:
@@ -625,6 +716,38 @@ class CodeSimilarityAnalyzer:
             by_class: dict[str, dict[tuple[str, str], dict[str, object]]] = {}
             for query_identifier, entry in result_map.items():
                 kind = entry.get("kind", "function")
+                if kind == "function":
+                    matches = entry.get("embedding", [])
+                    if not matches:
+                        continue
+                    best_per_cand: dict[tuple[str, str], dict[str, object]] = {}
+                    for match in matches:
+                        rel = match.get("relative_path")
+                        mname = match.get("match_name")
+                        score = match.get("score")
+                        if rel is None or not mname or score is None:
+                            continue
+                        if "." in mname:
+                            continue
+                        ckey = (rel, mname)
+                        prev = best_per_cand.get(ckey)
+                        if prev is None or float(score) > float(prev.get("score", -1.0)):
+                            best_per_cand[ckey] = match
+                    for ckey, match in best_per_cand.items():
+                        slot = by_class.setdefault(query_identifier, {}).setdefault(
+                            ckey,
+                            {
+                                "relative_path": ckey[0],
+                                "class_name": ckey[1],
+                                "scores": [],
+                                "contributors": [],
+                            },
+                        )
+                        slot["scores"].append(float(match["score"]))
+                        slot["contributors"].append(
+                            {"query": query_identifier, "match": match["identifier"], "score": float(match["score"])}
+                        )
+                    continue
                 qcls = query_class_key(query_identifier, kind)
                 matches = entry.get("embedding", [])
                 if not matches:
@@ -660,15 +783,26 @@ class CodeSimilarityAnalyzer:
             by_class_out: dict[str, list[dict[str, object]]] = {}
             for qcls, cand_map in by_class.items():
                 q_method_count = len(
-                    [key for key, kind in definitions_kind.items() if kind == "method" and key.startswith(f"{qcls}.")]
+                    [
+                        key
+                        for key, kind in definitions_kind.items()
+                        if kind == "method"
+                        and key.startswith(f"{qcls}.")
+                        and key.split(".")[-1] not in BOILERPLATE_NAMES
+                    ]
                 )
                 q_method_count = max(1, q_method_count)
                 rows = []
                 for _, slot in cand_map.items():
-                    scores = sorted(slot["scores"], reverse=True)
-                    coverage_count = len(scores)
+                    filtered_contributors = [
+                        item
+                        for item in slot["contributors"]
+                        if str(item.get("query", "")).split(".")[-1] not in BOILERPLATE_NAMES
+                    ]
+                    base_score, coverage_count = _calculate_reconstruction_score(
+                        filtered_contributors, q_method_count
+                    )
                     coverage_ratio = coverage_count / float(q_method_count)
-                    base_score = _calculate_reconstruction_score(scores, coverage_ratio)
                     contributors = sorted(slot["contributors"], key=lambda x: float(x["score"]), reverse=True)[:5]
                     rows.append(
                         {

app/main.py

@@ -157,12 +157,26 @@ def _get_structural_flow(node: ast.AST) -> str:
     flow: list[str] = []
     for child in ast.walk(node):
         if isinstance(child, ast.Call):
-            name = _call_name(child.func)
+            name = None
+            if isinstance(child.func, ast.Attribute) and isinstance(child.func.value, ast.Name):
+                if child.func.value.id == "self":
+                    name = f"self.{child.func.attr}"
+            if name is None:
+                name = _call_name(child.func)
             if name:
                 flow.append(name)
+        elif isinstance(child, ast.Attribute):
+            if isinstance(child.value, ast.Name) and child.value.id == "self":
+                flow.append(f"self.{child.attr}")
         elif isinstance(child, (ast.If, ast.While, ast.For)):
-            flow.append(f"Control({type(child).__name__})")
-    return " -> ".join(flow[:15])
+            flow.append("[LOGIC]")
+        elif isinstance(child, ast.Return):
+            flow.append("Return")
+    reduced: list[str] = []
+    for item in flow:
+        if not reduced or reduced[-1] != item:
+            reduced.append(item)
+    return " -> ".join(reduced[:20])
 
 
 def _extract_ast(source: str, symbol: str) -> tuple[str | None, dict[str, object] | None]:

static/app.js

@@ -162,6 +162,63 @@ function formatSummary(summary) {
   return parts.join(" · ") || "No structural summary.";
 }
 
+function escapeHtml(text) {
+  if (!text) return "";
+  return text
+    .replace(/&/g, "&")
+    .replace(/</g, "&lt;")
+    .replace(/>/g, "&gt;")
+    .replace(/\"/g, "&quot;")
+    .replace(/'/g, "&#039;");
+}
+
+function diffLCS(text1, text2) {
+  const lines1 = text1 ? text1.split("\n") : [];
+  const lines2 = text2 ? text2.split("\n") : [];
+  const n = lines1.length;
+  const m = lines2.length;
+  const dp = Array.from({ length: n + 1 }, () => Array(m + 1).fill(0));
+
+  for (let i = 1; i <= n; i += 1) {
+    for (let j = 1; j <= m; j += 1) {
+      if (lines1[i - 1] === lines2[j - 1]) {
+        dp[i][j] = dp[i - 1][j - 1] + 1;
+      } else {
+        dp[i][j] = Math.max(dp[i - 1][j], dp[i][j - 1]);
+      }
+    }
+  }
+
+  let i = n;
+  let j = m;
+  const result = [];
+  while (i > 0 || j > 0) {
+    if (i > 0 && j > 0 && lines1[i - 1] === lines2[j - 1]) {
+      result.unshift({ type: "same", text: lines1[i - 1] });
+      i -= 1;
+      j -= 1;
+    } else if (j > 0 && (i === 0 || dp[i][j - 1] >= dp[i - 1][j])) {
+      result.unshift({ type: "add", text: lines2[j - 1] });
+      j -= 1;
+    } else {
+      result.unshift({ type: "del", text: lines1[i - 1] });
+      i -= 1;
+    }
+  }
+  return result;
+}
+
+function renderDiff(text1, text2) {
+  const diff = diffLCS(text1, text2);
+  return diff
+    .map((part) => {
+      const cls = part.type === "add" ? "diff-add" : part.type === "del" ? "diff-del" : "";
+      const safe = escapeHtml(part.text || " ");
+      return `<div class="diff-row ${cls}"><div class="code-line">${safe}</div></div>`;
+    })
+    .join("");
+}
+
 function setAst(queryAst, matchAst, querySummary, matchSummary) {
   if (astQueryEl) {
     astQueryEl.textContent = queryAst || "AST not found.";
@@ -180,6 +237,26 @@ function setAst(queryAst, matchAst, querySummary, matchSummary) {
     const matchFlow = matchSummary?.flow ? matchSummary.flow : "unavailable";
     flowComparisonEl.textContent = `Selected flow:\n${queryFlow}\n\nMatch flow:\n${matchFlow}`;
   }
+
+  const diffContainer = document.querySelector(".code-diff-view");
+  if (diffContainer) {
+    diffContainer.innerHTML = "";
+    if (queryAst && matchAst) {
+      const diffWrapper = document.createElement("div");
+      diffWrapper.className = "diff-wrapper";
+      diffWrapper.innerHTML = renderDiff(queryAst, matchAst);
+      diffContainer.appendChild(diffWrapper);
+    } else {
+      const left = document.createElement("pre");
+      left.className = "code-block";
+      left.textContent = queryAst || "";
+      const right = document.createElement("pre");
+      right.className = "code-block";
+      right.textContent = matchAst || "";
+      diffContainer.appendChild(left);
+      diffContainer.appendChild(right);
+    }
+  }
 }
 
 async function loadAst(symbol, matchIdentifier) {
@@ -287,7 +364,7 @@ function renderBlueprint(byClass) {
     item.addEventListener("click", () => {
       document.querySelectorAll(".blueprint-item").forEach((el) => el.classList.remove("is-active"));
       item.classList.add("is-active");
-      const matchIdentifier = top.top_contributors?.[0]?.match || top.identifier;
+      const matchIdentifier = top.identifier;
       const activeName = document.getElementById("activeModuleName");
       if (activeName) {
         activeName.textContent = `${qcls} vs ${top.class_name}`;

static/styles.css

@@ -1,14 +1,14 @@
 @import url("https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@400;500;600;700&display=swap");
 
 :root {
-  --bg: #f6f0e6;
-  --panel: #fff7ee;
-  --ink: #1b1b1b;
-  --muted: #6b5f55;
-  --accent: #d6572b;
-  --accent-2: #2b6fd6;
-  --accent-3: #1b8d57;
-  --shadow: rgba(27, 27, 27, 0.1);
+  --bg: #f4f4f2;
+  --panel: #f9f9f7;
+  --ink: #1f1f1c;
+  --muted: #6f6b66;
+  --accent: #8d7b6e;
+  --accent-2: #8791a3;
+  --accent-3: #7e9a8b;
+  --shadow: rgba(31, 31, 28, 0.08);
 }
 
 * {
@@ -19,9 +19,9 @@ body {
   margin: 0;
   font-family: "Space Grotesk", system-ui, sans-serif;
   color: var(--ink);
-  background: radial-gradient(circle at 20% 20%, #ffe4c7 0%, transparent 55%),
-    radial-gradient(circle at 85% 15%, #f5d2e8 0%, transparent 40%),
-    radial-gradient(circle at 70% 80%, #d8f0e2 0%, transparent 45%),
+  background: radial-gradient(circle at 15% 20%, #e9e5df 0%, transparent 55%),
+    radial-gradient(circle at 80% 10%, #ece7e1 0%, transparent 40%),
+    radial-gradient(circle at 70% 80%, #e6e2dc 0%, transparent 45%),
     var(--bg);
 }
 
@@ -341,7 +341,7 @@ textarea {
 
 .code-diff-view {
   display: grid;
-  grid-template-columns: repeat(2, minmax(0, 1fr));
+  grid-template-columns: 1fr;
   gap: 16px;
 }
 
@@ -359,6 +359,37 @@ textarea {
   white-space: pre-wrap;
 }
 
+.diff-wrapper {
+  display: flex;
+  flex-direction: column;
+  background: #fff;
+  border: 1px solid #e3d6c8;
+  border-radius: 12px;
+  overflow: hidden;
+  font-family: "Space Grotesk", monospace;
+  font-size: 12px;
+}
+
+.diff-row {
+  display: block;
+  border-bottom: 1px solid #f0e8dc;
+}
+
+.code-line {
+  padding: 2px 8px;
+  white-space: pre;
+}
+
+.diff-del {
+  background-color: #ffe6e6;
+  text-decoration: line-through;
+  opacity: 0.7;
+}
+
+.diff-add {
+  background-color: #e6ffec;
+}
+
 .structural-notes {
   display: grid;
   grid-template-columns: repeat(2, minmax(0, 1fr));