Skip to content

Commit dd46890

Browse files
inakineitorinakineitor
authored
chore: extract evals shared utils (#210)
1 parent ac451e1 commit dd46890

4 files changed

Lines changed: 108 additions & 72 deletions

File tree

‎cli/commands/evals/paperqa_figures.py‎

Lines changed: 12 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -20,45 +20,22 @@
2020
import matplotlib.ticker as mticker
2121
import polars as pl
2222

23+
from .utils import (
24+
NODE_TYPE_LABELS,
25+
PALETTE,
26+
THEMES,
27+
load_polled_edges,
28+
normalize_relation_types,
29+
)
30+
2331
logger = logging.getLogger("cli")
2432

25-
# Human-readable labels for compact node-type codes
26-
NODE_TYPE_LABELS: dict[str, str] = {
27-
"ANA": "Anatomy",
28-
"BPO": "Biological process",
29-
"CCO": "Cellular component",
30-
"DIS": "Disease",
31-
"DRG": "Drug",
32-
"EXP": "Exposure",
33-
"GEN": "Gene",
34-
"MFN": "Molecular function",
35-
"PHE": "Phenotype",
36-
"PWY": "Pathway",
37-
}
38-
39-
_PALETTE = [
40-
"#516FD9", # Royal Blue
41-
"#7EACF5", # Sky Blue
42-
"#69C39C", # Mint
43-
"#6FA430", # Green
44-
"#E7C454", # Yellow
45-
"#EDB453", # Amber
46-
"#ED9353", # Orange
47-
"#DA3546", # Red
48-
"#9B7DF1", # Purple
49-
"#838E9F", # Gray
50-
]
33+
_PALETTE = PALETTE
34+
_THEMES = THEMES
5135

5236
_FALSE_GRAY = "#D0D0D0"
5337
_ALL_RATINGS = [1, 2, 3, 4, 5]
5438
_RATING_LABELS = ["No evidence", "Weak", "Moderate", "Strong", "Very strong"]
55-
56-
# Per-panel SVG themes. Data colors (palette, gray for false edges) stay the same;
57-
# only axes/text/spines flip between light and dark.
58-
_THEMES: dict[str, dict[str, str]] = {
59-
"light": {"ink": "#26251e", "muted": "#57534e"},
60-
"dark": {"ink": "#ebebeb", "muted": "#a8a29e"},
61-
}
6239
_RATING_NO_EVIDENCE = _ALL_RATINGS[0]
6340
_RATING_WEAK = _ALL_RATINGS[1]
6441
_RATING_MODERATE = _ALL_RATINGS[2]
@@ -72,41 +49,6 @@ def _run_id_from_path(path: Path) -> str:
7249
return path.stem
7350

7451

75-
def _load_df(input_path: Path) -> pl.DataFrame:
76-
if not input_path.exists():
77-
raise FileNotFoundError(f"Input file not found: {input_path}")
78-
79-
df = pl.read_csv(input_path, infer_schema_length=100000)
80-
81-
required = {"seed_node_type", "is_true_edge", "rating"}
82-
missing = required - set(df.columns)
83-
if missing:
84-
raise ValueError(f"Input CSV is missing required columns: {missing}")
85-
86-
return df.with_columns(
87-
pl.col("rating").cast(pl.Int32, strict=False),
88-
pl.col("is_true_edge").cast(pl.Boolean, strict=False),
89-
).filter(pl.col("rating").is_not_null() & pl.col("is_true_edge").is_not_null())
90-
91-
92-
# Whole-token PRO → GEN (Unicode word boundaries: hyphens/pipes/ends ok; not "PROGRAM")
93-
_RELATION_PRO_TO_GEN_PATTERN = r"\bPRO\b"
94-
95-
96-
def _relation_type_pro_to_gen(df: pl.DataFrame) -> pl.DataFrame:
97-
"""Replace whole-token ``PRO`` with ``GEN`` in ``relation_type`` (pipes, hyphens, etc.)."""
98-
if "relation_type" not in df.columns:
99-
return df
100-
return df.with_columns(
101-
pl.when(pl.col("relation_type").is_null())
102-
.then(None)
103-
.otherwise(
104-
pl.col("relation_type").str.replace_all(_RELATION_PRO_TO_GEN_PATTERN, "GEN")
105-
)
106-
.alias("relation_type"),
107-
)
108-
109-
11052
def _by_prevalence(df: pl.DataFrame, col: str) -> list[str]:
11153
"""Return unique values of `col` ordered by descending row count."""
11254
return (
@@ -643,12 +585,12 @@ def run(
643585
out_dir = out_dir or input_path.parent
644586
out_dir.mkdir(parents=True, exist_ok=True)
645587

646-
df = _load_df(input_path)
588+
df = load_polled_edges(input_path)
647589
if df.is_empty():
648590
logger.warning("No rows with valid rating + is_true_edge; nothing to plot.")
649591
return
650592

651-
df = _relation_type_pro_to_gen(df)
593+
df = normalize_relation_types(df)
652594

653595
run_id = _run_id_from_path(input_path)
654596

‎cli/commands/evals/utils.py‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,100 @@
1919
# Graph construction modes
2020
GraphMode = Literal["directed", "undirected"]
2121

22+
# Human-readable labels for compact node-type codes
23+
NODE_TYPE_LABELS: dict[str, str] = {
24+
"ANA": "Anatomy",
25+
"BPO": "Biological process",
26+
"CCO": "Cellular component",
27+
"DIS": "Disease",
28+
"DRG": "Drug",
29+
"EXP": "Exposure",
30+
"GEN": "Gene",
31+
"MFN": "Molecular function",
32+
"PHE": "Phenotype",
33+
"PWY": "Pathway",
34+
}
35+
36+
PALETTE = [
37+
"#516FD9", # Royal Blue
38+
"#7EACF5", # Sky Blue
39+
"#69C39C", # Mint
40+
"#6FA430", # Green
41+
"#E7C454", # Yellow
42+
"#EDB453", # Amber
43+
"#ED9353", # Orange
44+
"#DA3546", # Red
45+
"#9B7DF1", # Purple
46+
"#838E9F", # Gray
47+
]
48+
49+
# Per-panel SVG themes. Data colors stay fixed; only axes/text/spines flip.
50+
THEMES: dict[str, dict[str, str]] = {
51+
"light": {"ink": "#26251e", "muted": "#57534e"},
52+
"dark": {"ink": "#ebebeb", "muted": "#a8a29e"},
53+
}
54+
55+
# Whole-token PRO → GEN (Unicode word boundaries: hyphens/pipes/ends ok; not "PROGRAM")
56+
_RELATION_PRO_TO_GEN_PATTERN = r"\bPRO\b"
57+
58+
59+
def load_polled_edges(input_path: Path) -> pl.DataFrame:
60+
"""Load and normalize the polled-edges CSV produced by ``cli evals paperqa``.
61+
62+
Rows whose ``rating`` or ``is_true_edge`` could not be parsed are dropped;
63+
these correspond to queries for which the agent returned no usable answer.
64+
65+
Args:
66+
input_path: Path to the polled-edges CSV.
67+
68+
Returns:
69+
DataFrame with ``rating`` cast to Int32 and ``is_true_edge`` cast to
70+
Boolean, filtered to rows where both are non-null.
71+
72+
Raises:
73+
FileNotFoundError: If ``input_path`` does not exist.
74+
ValueError: If required columns are missing.
75+
"""
76+
if not input_path.exists():
77+
raise FileNotFoundError(f"Input file not found: {input_path}")
78+
79+
df = pl.read_csv(input_path, infer_schema_length=100000)
80+
81+
required = {"seed_node_type", "is_true_edge", "rating"}
82+
missing = required - set(df.columns)
83+
if missing:
84+
raise ValueError(f"Input CSV is missing required columns: {missing}")
85+
86+
return df.with_columns(
87+
pl.col("rating").cast(pl.Int32, strict=False),
88+
pl.col("is_true_edge").cast(pl.Boolean, strict=False),
89+
).filter(pl.col("rating").is_not_null() & pl.col("is_true_edge").is_not_null())
90+
91+
92+
def normalize_relation_types(df: pl.DataFrame) -> pl.DataFrame:
93+
"""Replace whole-token ``PRO`` with ``GEN`` in ``relation_type``.
94+
95+
OptimusKG is gene-centric: protein endpoints are represented as gene nodes,
96+
so relation labels recorded as ``PRO`` during sampling are reported as
97+
``GEN``. Returns the frame unchanged when ``relation_type`` is absent.
98+
99+
Args:
100+
df: DataFrame that may contain a ``relation_type`` column.
101+
102+
Returns:
103+
DataFrame with normalized relation-type labels.
104+
"""
105+
if "relation_type" not in df.columns:
106+
return df
107+
return df.with_columns(
108+
pl.when(pl.col("relation_type").is_null())
109+
.then(None)
110+
.otherwise(
111+
pl.col("relation_type").str.replace_all(_RELATION_PRO_TO_GEN_PATTERN, "GEN")
112+
)
113+
.alias("relation_type"),
114+
)
115+
22116

23117
def load_graph(
24118
nodes_path: Path,

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
66
requires-python = ">=3.12"
77
name = "optimuskg"
88
readme = "README.md"
9-
version = "0.72.5"
9+
version = "0.72.6"
1010
dependencies = [
1111
"ipython>=8.10",
1212
"jupyterlab>=3.0",

‎uv.lock‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)