mypy

Author: kddubey

.github/workflows/ci.yml

@@ -20,6 +20,17 @@ jobs:
       - run: uv tool run ruff check .
       - run: uv tool run ruff format --check .
 
+  typecheck:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+      - uses: astral-sh/setup-uv@eac588ad8def6316056a12d4907a9d4d84ff7a3b # v7.3.0
+        with:
+          enable-cache: true
+          python-version: "3.13"
+      - run: uv sync --extra dev
+      - run: uv run mypy
+
   test:
     runs-on: ubuntu-latest
     steps:

.pre-commit-config.yaml

@@ -1,5 +1,4 @@
 repos:
-  # Ruff lint + format. Config lives in pyproject.toml.
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.15.6
     hooks:
@@ -17,3 +16,13 @@ repos:
       - id: end-of-file-fixer
       - id: trailing-whitespace
       - id: debug-statements
+
+  - repo: local
+    hooks:
+      - id: mypy
+        name: mypy
+        entry: uv run mypy
+        language: system
+        types: [python]
+        pass_filenames: false
+        require_serial: true

.vscode/extensions.json

@@ -0,0 +1,6 @@
+{
+    "recommendations": [
+        "charliermarsh.ruff",
+        "ms-python.mypy-type-checker"
+    ]
+}

.vscode/settings.json

@@ -15,5 +15,8 @@
     "rewrap.autoWrap.enabled": true,
     "notebook.formatOnSave.enabled": true,
     "notebook.defaultFormatter": "charliermarsh.ruff",
-    "python.defaultInterpreterPath": ".venv/bin/python"
+    "python.defaultInterpreterPath": ".venv/bin/python",
+    "mypy-type-checker.reportingScope": "workspace",
+    "mypy-type-checker.importStrategy": "fromEnvironment",
+    "mypy-type-checker.preferDaemon": true
 }

benchmark/compare_compiled_uncompiled.py

@@ -25,6 +25,7 @@
 
 import numpy as np
 import polars as pl
+import torch
 from sentence_transformers.util import pairwise_cos_sim
 from tap import tapify
 
@@ -53,7 +54,7 @@ def _load_side(dir_local: Path) -> tuple[pl.DataFrame, np.ndarray, np.ndarray]:
 
 
 def _cos_sim_truncated(a: np.ndarray, b: np.ndarray, dim: int) -> np.ndarray:
-    return pairwise_cos_sim(a[..., :dim], b[..., :dim]).detach().cpu().numpy()
+    return pairwise_cos_sim(torch.as_tensor(a[..., :dim]), torch.as_tensor(b[..., :dim])).detach().cpu().numpy()
 
 
 def _percentiles(x: np.ndarray) -> dict[str, float]:
@@ -249,8 +250,8 @@ def main(
 
     n_pairs = len(df_joined)
 
-    def _fmt_count_pct(n: int) -> str:
-        return f"{n:,} ({n / n_pairs * 100:.3g}%)"
+    def _fmt_count_pct(n: int | float) -> str:
+        return f"{int(n):,} ({n / n_pairs * 100:.3g}%)"
 
     df_threshold = pl.DataFrame(
         {

benchmark/report.py

@@ -74,13 +74,13 @@ def _write_report(
     path_out: Path,
 ) -> None:
     speedup = df["time_base_sec"] / df["time_compiled_sec"]
-    p10 = float(speedup.quantile(0.1))
-    p50 = float(speedup.median())
-    p90 = float(speedup.quantile(0.9))
-    fraction_wins = float((speedup >= 1.0).mean())
+    p10 = float(speedup.quantile(0.1))  # type: ignore[arg-type]
+    p50 = float(speedup.median())  # type: ignore[arg-type]
+    p90 = float(speedup.quantile(0.9))  # type: ignore[arg-type]
+    fraction_wins = float((speedup >= 1.0).mean())  # type: ignore[arg-type]
 
-    median_compiled_ms = float(df["time_compiled_sec"].median()) * 1000
-    median_base_ms = float(df["time_base_sec"].median()) * 1000
+    median_compiled_ms = float(df["time_compiled_sec"].median()) * 1000  # type: ignore[arg-type]
+    median_base_ms = float(df["time_base_sec"].median()) * 1000  # type: ignore[arg-type]
 
     df_worst = (
         df.with_columns(speedup=pl.col("time_base_sec") / pl.col("time_compiled_sec"))

eval/compare.py

@@ -46,6 +46,7 @@
 from dataclasses import dataclass
 from itertools import zip_longest
 from pathlib import Path
+from typing import Any
 from unittest.mock import patch
 
 import gspread
@@ -313,7 +314,7 @@ def _compute_metrics_for_model(df: pl.DataFrame, model_name: str) -> dict:
     }
 
 
-def _compute_metrics_avg_over_projects(df: pl.DataFrame, model_name: str) -> dict:
+def _compute_metrics_avg_over_projects(df: pl.DataFrame, model_name: str) -> dict[str, float]:
     """Compute metrics averaged over projects so large projects don't dominate."""
     metrics_per_project = []
     for _, df_project in df.group_by("project_id"):
@@ -350,14 +351,16 @@ def plot_metrics_by_platform(df: pl.DataFrame, model_names: list[str]) -> plt.Fi
     metrics_to_plot = None
     for (platform,), platform_df in df.group_by("platform"):
         for model_name in model_names:
-            project_metrics_list = []
+            project_metrics_list: list[dict[str, Any]] = []
             for _, proj_df in platform_df.group_by("project_id"):
                 project_metrics_list.append(_compute_metrics_for_model(proj_df, model_name))
             if metrics_to_plot is None:
                 metrics_to_plot = list(project_metrics_list[0].keys())
             avg_metrics = {
-                k: sum(m[k] for m in project_metrics_list if m[k] == m[k])
-                / sum(1 for m in project_metrics_list if m[k] == m[k])
+                k: (
+                    sum(m[k] for m in project_metrics_list if m[k] == m[k])
+                    / sum(1 for m in project_metrics_list if m[k] == m[k])
+                )
                 for k in project_metrics_list[0]
             }
             avg_metrics["platform"] = platform
@@ -367,9 +370,10 @@ def plot_metrics_by_platform(df: pl.DataFrame, model_names: list[str]) -> plt.Fi
     metrics_df = pl.DataFrame(metrics_rows)
 
     # Convert to pandas and pivot for plotting
+    assert metrics_to_plot is not None, "No platforms in df"
     metrics_pd = metrics_df.to_pandas()
-    fig, axes = plt.subplots(1, len(metrics_to_plot), figsize=(4 * len(metrics_to_plot), 5))
-    axes: list[plt.Axes] = list(axes)
+    fig, axes_arr = plt.subplots(1, len(metrics_to_plot), figsize=(4 * len(metrics_to_plot), 5))
+    axes: list[plt.Axes] = list(axes_arr)
 
     for ax, metric in zip(axes, metrics_to_plot, strict=True):
         pivot_df = metrics_pd.pivot(index="platform", columns="model", values=metric)
@@ -382,7 +386,7 @@ def plot_metrics_by_platform(df: pl.DataFrame, model_names: list[str]) -> plt.Fi
     # Single legend for the whole figure (top center)
     handles, labels = axes[0].get_legend_handles_labels()
     fig.legend(handles, labels, loc="upper center", ncol=len(model_names), bbox_to_anchor=(0.5, 1.02))
-    plt.tight_layout(rect=[0, 0, 1, 0.95])  # make room for legend on top
+    plt.tight_layout(rect=(0, 0, 1, 0.95))  # make room for legend on top
     return fig
 
 
@@ -443,10 +447,8 @@ def plot_dumbbell_by_project(
         metrics = [c.replace(f"{model1}_", "") for c in project_metrics_df.columns if c.startswith(f"{model1}_")]
 
     n_metrics = len(metrics)
-    fig, axes = plt.subplots(1, n_metrics, figsize=(5 * n_metrics, max(8, len(project_metrics_df) * 0.15)))
-    if n_metrics == 1:
-        axes = [axes]
-    axes: list[plt.Axes] = list(axes)
+    fig, axes_arr = plt.subplots(1, n_metrics, figsize=(5 * n_metrics, max(8, len(project_metrics_df) * 0.15)))
+    axes: list[plt.Axes] = [axes_arr] if n_metrics == 1 else list(axes_arr)
 
     # Sort once by pred_GROUP_rate delta, use same order for all subplots
     group_rate_col1 = f"{model1}_pred_GROUP_rate"
@@ -483,7 +485,7 @@ def plot_dumbbell_by_project(
     handles, labels = axes[0].get_legend_handles_labels()
     fig.legend(handles, labels, loc="upper center", ncol=len(model_names), bbox_to_anchor=(0.5, 1.02))
     fig.suptitle("Metrics by Project (org_id|project_id)", fontsize=14, y=1.05)
-    plt.tight_layout(rect=[0, 0, 1, 0.98])
+    plt.tight_layout(rect=(0, 0, 1, 0.98))
     return fig
 
 
@@ -557,11 +559,13 @@ def compare_models(
     # Compute conditional probabilities (reported later)
     prod_group = df.filter(pl.col(pred1_col) == "GROUP")
     prod_separate = df.filter(pl.col(pred1_col) == "SEPARATE")
-    p_group_given_group = (prod_group[pred2_col] == "GROUP").mean() if len(prod_group) > 0 else float("nan")
-    p_group_given_separate = (prod_separate[pred2_col] == "GROUP").mean() if len(prod_separate) > 0 else float("nan")
+    p_group_given_group = float((prod_group[pred2_col] == "GROUP").mean()) if len(prod_group) > 0 else float("nan")  # type: ignore[arg-type]
+    p_group_given_separate = (
+        float((prod_separate[pred2_col] == "GROUP").mean()) if len(prod_separate) > 0 else float("nan")  # type: ignore[arg-type]
+    )
     df_close = df.filter(pl.col("distance") < 0.005)
     close_group = df_close.filter(pl.col(pred1_col) == "GROUP")
-    p_close = (close_group[pred2_col] == "GROUP").mean() if len(close_group) > 0 else float("nan")
+    p_close = float((close_group[pred2_col] == "GROUP").mean()) if len(close_group) > 0 else float("nan")  # type: ignore[arg-type]
 
     # Columns to keep in output
     output_cols = [
@@ -589,6 +593,7 @@ def compare_models(
     df_sorted = df.sort(["org_id", "project_id"])
     for (org_id, project_id), group_df in df_sorted.group_by(["org_id", "project_id"], maintain_order=True):
         total_projects += 1
+        assert output_dir is not None
         proj_dir = output_dir / f"org_{org_id}" / f"project_{project_id}"
 
         # Compute metrics for each model on this project
@@ -675,7 +680,7 @@ def compare_models(
 
     report("\n### Distance distribution\n")
     report(df["distance"].describe())
-    report(f"\nGROUP rate: {(df['label'] == 'GROUP').mean():.2%}")
+    report(f"\nGROUP rate: {float((df['label'] == 'GROUP').mean()):.2%}")  # type: ignore[arg-type]
 
     platform_stats = (
         df.group_by("platform")
@@ -754,7 +759,7 @@ def compute_stacktrace_token_percentiles(df: pl.DataFrame) -> pl.DataFrame:
 
     rows = []
     for col in token_cols:
-        row = {"metric": col}
+        row: dict[str, Any] = {"metric": col}
         row["min"] = df[col].min()
         row["mean"] = df[col].mean()
         for p in percentiles:
@@ -896,6 +901,7 @@ def _compute_project_precisions_per_platform(model: str, thresholds_platform: di
             )
         else:
             baseline_key = f"{baseline_model}@{baseline_threshold}"
+            assert isinstance(baseline_threshold, float)
             project_precisions[baseline_key] = _compute_project_precisions(baseline_model, baseline_threshold)
     else:
         baseline_key = str(thresholds_sorted[0])
@@ -916,7 +922,7 @@ def _compute_project_precisions_per_platform(model: str, thresholds_platform: di
                 {
                     "platform": platform,
                     "n_projects": len(prec),
-                    "median_pairs": int(platform_df["n_pairs"].median()),
+                    "median_pairs": int(platform_df["n_pairs"].median()),  # type: ignore[arg-type]
                     "mean": prec.mean(),
                     "p5": prec.quantile(0.05),
                     "p10": prec.quantile(0.10),
@@ -971,7 +977,8 @@ def metrics_by_platform(
         )
 
     rows = []
-    for (platform,), platform_df in df_t.group_by("platform"):
+    for (platform_obj,), platform_df in df_t.group_by("platform"):
+        platform = str(platform_obj)
         avg_metrics = _compute_metrics_avg_over_projects(platform_df, model_name)
         platform_threshold = threshold.get(platform, threshold["default"]) if isinstance(threshold, dict) else threshold
         rows.append(
@@ -1025,12 +1032,15 @@ def find_threshold_by_platform(
     precision_by_platform = min_precision if isinstance(min_precision, dict) else None
 
     rows = []
-    for (platform,), platform_df in df.group_by("platform"):
+    for (platform_obj,), platform_df in df.group_by("platform"):
+        platform = str(platform_obj)
         n_pairs = len(platform_df)
         n_projects = platform_df["project_id"].n_unique()
         label_group_rate = (platform_df["label"] == "GROUP").mean()
         threshold_found = None
-        target_precision = precision_by_platform[platform] if precision_by_platform else min_precision
+        target_precision: float = (
+            precision_by_platform[platform] if precision_by_platform else min_precision  # type: ignore[assignment]
+        )
 
         # Walk thresholds from low to high; first one meeting precision is the minimum
         # Precision is averaged over projects to avoid large projects dominating
@@ -1039,11 +1049,11 @@ def find_threshold_by_platform(
                 pl.when(pl.col(sim_col) > thresh).then(pl.lit("GROUP")).otherwise(pl.lit("SEPARATE")).alias(pred_col)
             )
             # Compute per-project precision, then average
-            project_precisions = []
+            project_precisions: list[float] = []
             for _, proj_df in df_t.group_by("project_id"):
                 pred_group = proj_df.filter(pl.col(pred_col) == "GROUP")
                 if len(pred_group) > 0:
-                    project_precisions.append((pred_group["label"] == "GROUP").mean())
+                    project_precisions.append(float((pred_group["label"] == "GROUP").mean()))  # type: ignore[arg-type]
             if not project_precisions:
                 continue
             precision = sum(project_precisions) / len(project_precisions)
@@ -1115,11 +1125,11 @@ def compare_metrics_by_stacktrace_length(
 
     # Print metrics for each bucket
     report(f"\n### Short stacktraces ({token_col} <= p10 = {p10:.0f} tokens, {len(short_df)} pairs)\n")
-    report(f"label GROUP rate: {(short_df['label'] == 'GROUP').mean():.2%}")
+    report(f"label GROUP rate: {float((short_df['label'] == 'GROUP').mean()):.2%}")  # type: ignore[arg-type]
     report(_compute_metrics(short_df, model_names))
 
     report(f"\n### Long stacktraces ({token_col} >= p90 = {p90:.0f} tokens, {len(long_df)} pairs)\n")
-    report(f"label GROUP rate: {(long_df['label'] == 'GROUP').mean():.2%}")
+    report(f"label GROUP rate: {float((long_df['label'] == 'GROUP').mean()):.2%}")  # type: ignore[arg-type]
     report(_compute_metrics(long_df, model_names))
 
 

eval/eval_poller.py

@@ -154,6 +154,7 @@ def evaluate_baseline(
         return
 
     logger.info("Evaluating base model.")
+    loss: gt.loss.PairwiseLoss
     if loss_type == "sigmoid":
         loss = gt.loss.SigmoidPairwiseLoss()
     elif loss_type == "contrastive":

eval/save_embeddings.py

@@ -178,28 +178,36 @@ def main(
             text_prefix=text_prefix,
         )
         logger.info(f"{st_class.__name__} loaded in {time.monotonic() - start:.1f}s")
-        if use_compiled:
+        if isinstance(model, gt.compiled.SentenceTransformer):
             model.compile_and_warm_up()
         else:
             _ = model.encode("warm up")
         logger.info(f"{st_class.__name__} loading and warming up done in {time.monotonic() - start:.1f}s")
 
         logger.info("Encoding queries")
         texts_query = df["query_stacktrace_string"].to_list()
-        embeddings_query: np.ndarray = model.encode(
+        embeddings_query = model.encode(
             texts_query, batch_size=batch_size, convert_to_numpy=True, show_progress_bar=True
         )
         logger.info("Encoding candidates")
         texts_candidate = df["candidate_stacktrace_string"].to_list()
-        embeddings_candidate: np.ndarray = model.encode(
+        embeddings_candidate = model.encode(
             texts_candidate, batch_size=batch_size, convert_to_numpy=True, show_progress_bar=True
         )
 
     if truncate_dims is None:
         truncate_dims = (embeddings_query.shape[-1],)
 
     for dim in truncate_dims:
-        cos_sims = pairwise_cos_sim(embeddings_query[..., :dim], embeddings_candidate[..., :dim]).detach().cpu().numpy()
+        cos_sims = (
+            pairwise_cos_sim(
+                torch.as_tensor(embeddings_query[..., :dim]),
+                torch.as_tensor(embeddings_candidate[..., :dim]),
+            )
+            .detach()
+            .cpu()
+            .numpy()
+        )
         df = df.with_columns(pl.Series(name=f"cos_sim_{dim}", values=cos_sims))
 
     with tempfile.TemporaryDirectory() as dir_tmp_output:

eval/save_gemini_embeddings.py

@@ -20,6 +20,7 @@
 
 import numpy as np
 import polars as pl
+import torch
 from google import genai
 from google.genai import types
 from sentence_transformers.util import pairwise_cos_sim
@@ -148,7 +149,15 @@ def main(
     embeddings_candidate = all_embeddings[n_q:]
 
     for dim in truncate_dims:
-        cos_sims = pairwise_cos_sim(embeddings_query[..., :dim], embeddings_candidate[..., :dim]).detach().cpu().numpy()
+        cos_sims = (
+            pairwise_cos_sim(
+                torch.as_tensor(embeddings_query[..., :dim]),
+                torch.as_tensor(embeddings_candidate[..., :dim]),
+            )
+            .detach()
+            .cpu()
+            .numpy()
+        )
         df = df.with_columns(pl.Series(name=f"cos_sim_{dim}", values=cos_sims))
 
     with tempfile.TemporaryDirectory() as dir_tmp_output: