Skip to content

Refactor/Cleanup - #5

Merged
dimalvovs merged 10 commits into
mainfrom
refactor-cleanup
Aug 3, 2026
Merged

Refactor/Cleanup#5
dimalvovs merged 10 commits into
mainfrom
refactor-cleanup

Conversation

@MateoZ05

Copy link
Copy Markdown
Collaborator

Cleaned up llmize entire codebase for shipping, removes tooluniverse and enrichment stack, consolidate JSON-reduction modules, and fixes stalls with smaller models.

What changed

Removals / decoupling

Deleted enrich.py and all ToolUniverse/enrichment wiring (--enrich, deps, Docker/Nextflow flags); moved extract_entities into verify.py (its only consumer).
Removed dead code: call_ollama, json_reduction/main.py, json_reduction/json_topKeys.py, plus unused imports/params. Dropped tooluniverse/PyYAML from requirements and the Docker build.

Simplification

Unified whole-report and section-by-section into a single interpret_report() (whole report = one chunk; synthesis only when >1).
interpret.py is now a pure library; pipeline.py is the single CLI entry point.
Consolidated json_reduction/ (json_load/clean/write/merger → one reduction.py).
Reduction + annotation now run in memory; --save-intermediates writes the intermediate JSON only when requested.

Fixes

Fixed a small-model stall on numeric sections (e.g. deconvolved_probs): _round_floats() strips spurious 17-digit float precision from prompts, which was triggering repetition loops.
Reconciled with main: preserved the gene-split (verify.py), prompt-injection escaping (main.nf), prompt tweaks, and the data-grounding review feature; repaired the review path's extract_entities import after enrich removal.

Verification
Deterministic outputs (annotated report, entities, prompts) byte-identical to the pre-refactor baseline, except the intended prompt-float rounding.
compileall and pipeline.py --check pass; no reintroduced enrich/ToolUniverse references.
Full offline Docker and Nextflow runs complete end-to-end with a small model.
Summary generated by Claude Code

@MateoZ05
MateoZ05 requested a review from dimalvovs July 21, 2026 21:57

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR refactors llmize toward a shippable, CLI-driven pipeline by removing the ToolUniverse/enrichment stack, consolidating JSON reduction/annotation into a single in-memory API, and simplifying interpretation into a library function invoked by pipeline.py.

Changes:

  • Removes enrichment/ToolUniverse wiring (CLI flags, checks, Docker/Nextflow plumbing) and deletes enrich.py.
  • Consolidates json_reduction/* helpers into json_reduction/reduction.py and shifts the pipeline to run reduction+annotation in memory (optional --save-intermediates).
  • Unifies interpretation into interpret_report() with float-rounding to reduce small-model digit-echo stalls.

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
verify.py Moves entity extraction logic into verify.py for the review path and post-enrichment cleanup.
pipeline.py Becomes the single CLI entrypoint; runs reduction+annotation in memory and calls interpret_report().
interpret.py Refactors interpretation into a library API, adds float rounding for prompts, and updates review integration.
json_reduction/reduction.py New consolidated API for resolving/loading/saving/reducing/annotating MultiQC JSON.
json_reduction/init.py Re-exports the new consolidated reduction/annotation API.
nextflow.config Removes the enrich parameter.
main.nf Removes --enrich flag wiring from the Nextflow process invocation.
check_env.py Removes ToolUniverse/enrichment checks from environment diagnostics.
Dockerfile Drops enrichment-related install flags/deps; installs only ollama.
docker/entrypoint.sh Minor comment cleanup; keeps default --check behavior.
docker/boot_ollama.sh Comment cleanup; keeps offline-safe model pull behavior.
enrich.py Deleted (ToolUniverse enrichment stack removed).
json_reduction/main.py Deleted (interactive reduction entrypoint removed).
json_reduction/json_load.py Deleted (superseded by reduction.py).
json_reduction/json_clean.py Deleted (superseded by reduction.py).
json_reduction/json_write.py Deleted (superseded by reduction.py).
json_reduction/json_topKeys.py Deleted (dead code removed).

Comment thread interpret.py Outdated
Comment on lines 758 to 762
@@ -775,135 +759,4 @@ def g(key):
return "\n".join(lines)


def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Pass an annotated MultiQC JSON report to a local Ollama model."
)
parser.add_argument(
"--input", "-i",
required=True,
help="Path to the annotated JSON report (output of main.py)",
)
parser.add_argument(
"--model", "-m",
default="gemma4",
help="Ollama model name to use (default: llama3)",
)
parser.add_argument(
"--output", "-o",
default=None,
help="Path to save the LLM response. Defaults to <input_stem>_interpretation.txt",
)
parser.add_argument(
"--num_ctx",
type=int,
default=32768,
help="Context window size (default: 32768)",
)
parser.add_argument(
"--whole-report",
action="store_true",
help="Send the entire report in a single call instead of analysing each section one-by-one.",
)
parser.add_argument(
"--enrich",
action="store_true",
help="Enrich the system prompt with ToolUniverse gene annotations (requires the 'tooluniverse' package).",
)
parser.add_argument(
"--no-synthesis",
action="store_true",
help="Skip the final executive-summary synthesis pass (section-by-section mode only).",
)
parser.add_argument(
"--think",
action=argparse.BooleanOptionalAction,
default=True,
help="Use the model's native thinking mode; the reasoning is printed to the terminal "
"(not saved) and kept out of the interpretation. On by default; --no-think disables it (faster).",
)
parser.add_argument(
"--prompt",
default=None,
help="Extra instruction appended to the system prompt for every section "
"(e.g. 'Focus on immune cell types' or 'Answer in bullet points').",
)
parser.add_argument("--temperature", type=float, default=None,
help="Sampling temperature (model default if unset; gemma4 defaults to 1).")
parser.add_argument("--top_p", type=float, default=None, help="Nucleus sampling top_p.")
parser.add_argument("--top_k", type=int, default=None, help="Top-k sampling.")
parser.add_argument("--seed", type=int, default=None,
help="Sampling seed for reproducible runs.")
parser.add_argument("--num_predict", type=int, default=None,
help="Max tokens to generate per call (model default if unset).")
return parser.parse_args()


def _build_enrichment_context(report: dict, enabled: bool) -> str:
if not enabled:
return ""
try:
from enrich import enrich_report
except Exception as exc:
print(f"[interpret] Enrichment unavailable ({exc}); continuing without it.")
return ""
return enrich_report(report)


def main() -> None:
args = parse_args()

if args.output is None:
stem = os.path.splitext(os.path.basename(args.input))[0]
out_dir = os.path.dirname(os.path.abspath(args.input))
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
args.output = os.path.join(out_dir, f"{stem}_interpretation_{timestamp}.md")

print(f"[interpret] Loading report: {args.input}")
report = load_report(args.input)

section_count = sum(1 for v in report.values() if isinstance(v, dict) and "data" in v)
print(f"[interpret] Loaded {section_count} data sections.")

enrichment = _build_enrichment_context(report, args.enrich)
gen_options = build_gen_options(
temperature=args.temperature, top_p=args.top_p, top_k=args.top_k,
seed=args.seed, num_predict=args.num_predict,
)

if args.whole_report:
prompt = build_prompt(report)
system = (SYSTEM_PROMPT + enrichment + build_glossary_context(report)
+ build_user_instruction(args.prompt))
response_text, thinking = call_ollama(
prompt, model=args.model, system=system,
num_ctx=args.num_ctx, think=args.think, gen_options=gen_options,
)
print_thinking("Whole report", thinking)
else:
response_text = interpret_per_section(
report,
model=args.model,
num_ctx=args.num_ctx,
extra_system_context=enrichment,
synthesize_final=not args.no_synthesis,
think=args.think,
gen_options=gen_options,
user_instruction=args.prompt,
)

footer = build_run_footer(
model=args.model, num_ctx=args.num_ctx, think=args.think, gen_options=gen_options,
mode="whole-report" if args.whole_report else "section-by-section",
enrich=args.enrich, user_instruction=args.prompt, input_path=args.input,
)
save_output(response_text + footer, args.output)

print("LLM RESPONSE PREVIEW (first 1000 chars)")
print(response_text[:1000])
if len(response_text) > 1000:
print(f"\n... [{len(response_text) - 1000} more characters in saved file]")


if __name__ == "__main__":
main() No newline at end of file
Comment thread interpret.py

summary = None
if synthesize_final and responses:
if synthesize_final:
Comment thread verify.py
Comment on lines +55 to +67
def _parse_ligrec_key(key: str, cell_types: set) -> dict | None:
remainder = key
trailing = []
for _ in range(2):
match = None
for ct in cell_types:
suffix = "-" + ct
if remainder.endswith(suffix) and (match is None or len(ct) > len(match)):
match = ct
if match is None:
break
trailing.insert(0, match)
remainder = remainder[: -(len(match) + 1)]
@dimalvovs
dimalvovs merged commit e9d62cd into main Aug 3, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants