Spaces:
Running
Running
MeluXina user
taxonova 26.11.3: short reads (80 to 119 valid bases) are classified, with a note in the results and a hint before the run; pin 26.11.3
4ffc1bd | """The classification pipeline the web app runs: ``taxonova classify`` itself. | |
| Each request's sequences are written to FASTA and classified by the package's | |
| own ``classification_mode`` with a model this process already holds (see | |
| ``package_adapter``), so the TSV is the one ``taxonova classify`` writes for | |
| the same model, input and thresholds. The rows of that TSV are what the user | |
| downloads. | |
| Thresholds: with no confidence and no per-rank threshold the model's | |
| calibrated operating point is used, as on the command line; any threshold | |
| switches to the scaled profile (``taxonova.core.thresholds``). The species | |
| gate and the species-group rule of a model apply in both cases, as on the | |
| command line. | |
| Sequences that fail the SSU rRNA screen (``rrna_prefilter``) never reach the | |
| model and are reported as Unclassified at every rank. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import os | |
| import tempfile | |
| from typing import Callable, Dict, List, Optional, Sequence, Tuple | |
| from taxonova.core.thresholds import GRADED_SPECIES | |
| from webapp import package_adapter | |
| from webapp.package_adapter import LoadedModel, RunReport, threshold_argv | |
| from webapp.rrna_prefilter import prefilter_sequences | |
| logger = logging.getLogger(__name__) | |
| #: Canonical rank order and TSV columns of the CLI. | |
| LEVEL_ORDER = list(package_adapter.RANKS) | |
| #: The species confidence of the legacy /classify endpoint when none is sent | |
| #: (its 26.10 default, kept so that old scripts get the same thresholds). | |
| #: The page and /api/jobs send no confidence by default: the model's | |
| #: operating point. | |
| LEGACY_CONFIDENCE = GRADED_SPECIES | |
| #: Sequences classified per model call. Labels do not depend on it (every | |
| #: sequence is classified independently); it sets how often progress is | |
| #: reported. test_parity.py checks that. A multiple of the classifier batch | |
| #: size, so batches line up with the command line. | |
| CHUNK_SIZE = int(os.environ.get("TAXONOVA_WEB_CHUNK", "1024")) | |
| Triple = Tuple[str, str, dict] | |
| def load_model(model_path: str) -> LoadedModel: | |
| """Load one checkpoint with the package loader.""" | |
| return package_adapter.load_model(model_path) | |
| def cli_command(model_id: str, confidence: Optional[float] = None, | |
| overrides: Optional[Dict[str, float]] = None) -> str: | |
| """The command that reproduces a web run on the command line.""" | |
| argv = threshold_argv(confidence, overrides) | |
| shown = [a if a.startswith("--") else f"{float(a):g}" for a in argv] | |
| return " ".join(["taxonova classify reads.fasta --model", model_id] + shown) | |
| def run_report(cached: LoadedModel, sequences: Sequence[Triple], confidence: Optional[float] = None, | |
| overrides: Optional[Dict[str, float]] = None, strand_check: bool = True) -> RunReport: | |
| """Classify ``sequences`` with a loaded model: the CLI TSV and what was withheld. | |
| ``strand_check`` False skips the package's orientation check (see | |
| ``package_adapter.classify``).""" | |
| tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".fasta", delete=False, encoding="utf-8") | |
| try: | |
| for seq_id, seq, _ in sequences: | |
| tmp.write(f">{seq_id}\n{seq}\n") | |
| tmp.close() | |
| return package_adapter.classify(cached, tmp.name, confidence, overrides, strand_check) | |
| finally: | |
| try: | |
| os.unlink(tmp.name) | |
| except OSError: | |
| pass | |
| def run_classification(cached: LoadedModel, sequences: Sequence[Triple], confidence: Optional[float] = None, | |
| overrides: Optional[Dict[str, float]] = None) -> str: | |
| """Classify ``sequences`` with a loaded model; return the CLI TSV text.""" | |
| return run_report(cached, sequences, confidence, overrides).tsv | |
| def tsv_header() -> str: | |
| cols = ["sequence_id"] | |
| for level in LEVEL_ORDER: | |
| cols += [level, f"{level}_confidence"] | |
| return "\t".join(cols) | |
| def unclassified_tsv_row(seq_id: str) -> str: | |
| return "\t".join([seq_id] + ["Unclassified", "0.0000"] * len(LEVEL_ORDER)) | |
| def uniquify_sequence_ids(sequences): | |
| """Give every sequence a unique internal id, keeping the display names. | |
| FASTA headers need not be unique and the merge keys on the id, so | |
| duplicates would otherwise all receive one call. | |
| Returns (sequences_with_unique_ids, display_names_by_position). | |
| """ | |
| used = set() | |
| unique, display = [], [] | |
| for seq_id, seq, tax in sequences: | |
| display.append(seq_id) | |
| uid, n = seq_id, 1 | |
| while uid in used: | |
| n += 1 | |
| uid = f"{seq_id}__dup{n}" | |
| used.add(uid) | |
| unique.append((uid, seq, tax)) | |
| return unique, display | |
| def restore_display_ids(tsv_string, results, display): | |
| """Put the user's sequence names back, by position.""" | |
| if len(results) != len(display): | |
| raise RuntimeError(f"result/display length mismatch: {len(results)} vs {len(display)}") | |
| for entry, name in zip(results, display): | |
| entry["sequence_id"] = name | |
| lines = tsv_string.splitlines() | |
| if not lines: | |
| return tsv_string, results | |
| body = lines[1:] | |
| if len(body) != len(display): | |
| raise RuntimeError(f"TSV/display length mismatch: {len(body)} vs {len(display)}") | |
| out = [lines[0]] | |
| for ln, name in zip(body, display): | |
| cols = ln.split("\t") | |
| cols[0] = name | |
| out.append("\t".join(cols)) | |
| return "\n".join(out) + "\n", results | |
| def parse_tsv(tsv_string: str) -> List[dict]: | |
| """CLI TSV text -> one dict per row (confidences as floats).""" | |
| lines = [ln for ln in tsv_string.splitlines() if ln != ""] | |
| if not lines: | |
| return [] | |
| header = lines[0].split("\t") | |
| col = {name: i for i, name in enumerate(header)} | |
| rows = [] | |
| for line in lines[1:]: | |
| f = line.split("\t") | |
| entry = {"sequence_id": f[col["sequence_id"]]} | |
| for level in LEVEL_ORDER: | |
| entry[level] = f[col[level]] if level in col else "Unclassified" | |
| try: | |
| entry[f"{level}_confidence"] = float(f[col[f"{level}_confidence"]]) | |
| except (KeyError, ValueError, IndexError): | |
| entry[f"{level}_confidence"] = 0.0 | |
| rows.append(entry) | |
| return rows | |
| ProgressFn = Callable[[str, float], None] | |
| def classify_with_prefilter(cached: LoadedModel, sequences: Sequence[Triple], confidence: Optional[float] = None, | |
| overrides: Optional[Dict[str, float]] = None, | |
| progress: Optional[ProgressFn] = None, | |
| chunk_size: int = CHUNK_SIZE) -> Tuple[str, List[dict], dict]: | |
| """SSU screen, then classify what passed, then merge in input order. | |
| ``sequences`` must already carry unique ids. Returns the TSV text (CLI | |
| schema; passing rows byte-identical to the CLI), one dict per input | |
| sequence with a ``prefilter_pass`` flag, and the run summary | |
| (``run_info``). | |
| """ | |
| report = progress or (lambda stage, frac: None) | |
| report("Screening for 16S rRNA", 0.02) | |
| pass_ids = prefilter_sequences(list(sequences)) | |
| passing = [t for t in sequences if t[0] in pass_ids] | |
| logger.info("rRNA screen: %d of %d passed", len(passing), len(sequences)) | |
| header = tsv_header() | |
| row_by_id: Dict[str, str] = {} | |
| reports: List[RunReport] = [] | |
| done = 0 | |
| for start, stop in chunk_bounds(len(passing), chunk_size): | |
| chunk = passing[start:stop] | |
| report(f"Classifying sequences ({done:,} of {len(passing):,})", | |
| 0.15 + 0.8 * done / max(1, len(passing))) | |
| # orientation: checked once per job, on the first reads, as the CLI checks a file | |
| rep = run_report(cached, chunk, confidence, overrides, strand_check=start == 0) | |
| reports.append(rep) | |
| lines = [ln for ln in rep.tsv.splitlines() if ln] | |
| if lines: | |
| header = lines[0] | |
| sid_idx = header.split("\t").index("sequence_id") | |
| for ln in lines[1:]: | |
| row_by_id[ln.split("\t")[sid_idx]] = ln | |
| done += len(chunk) | |
| report("Preparing results", 0.97) | |
| tsv_lines = [header] | |
| for seq_id, _, _ in sequences: | |
| tsv_lines.append(row_by_id.get(seq_id) or unclassified_tsv_row(seq_id)) | |
| tsv = "\n".join(tsv_lines) + "\n" | |
| rows = parse_tsv(tsv) | |
| for row, (seq_id, _, _) in zip(rows, sequences): | |
| row["prefilter_pass"] = seq_id in row_by_id | |
| return tsv, rows, run_info(cached, reports, passing, confidence, overrides) | |
| def chunk_bounds(n: int, chunk_size: int) -> List[Tuple[int, int]]: | |
| """(start, stop) of each chunk. The first chunk is made of whole chunks | |
| enough to hold the reads the package checks for orientation, so that the | |
| one check per job sees the same reads as the command line.""" | |
| size = max(1, chunk_size) | |
| sample = package_adapter.strand_sample_size() | |
| first = min(n, -(-sample // size) * size) | |
| bounds = [(0, first)] if first else [] | |
| bounds += [(s, min(s + size, n)) for s in range(first, n, size)] | |
| return bounds | |
| def run_info(cached: LoadedModel, reports: Sequence[RunReport], passing: Sequence[Triple], | |
| confidence: Optional[float], overrides: Optional[Dict[str, float]]) -> dict: | |
| """What the page says about a run: thresholds and their source, and why species are missing.""" | |
| settings = getattr(cached, "settings", None) or {} | |
| first = reports[0] if reports else None | |
| if first is not None: | |
| thresholds, source = first.thresholds, first.source | |
| gate, groups = first.species_gate, first.species_groups | |
| else: # nothing passed the screen: what would have applied | |
| thresholds = None | |
| source = "user" if (confidence is not None or overrides) else ( | |
| "checkpoint" if settings.get("operating_point") else "default") | |
| gate, groups = bool(settings.get("species_gate")), settings.get("species_groups", "kept") | |
| return { | |
| "thresholds": thresholds, | |
| "source": source, | |
| "species_gate": gate, | |
| "species_groups": groups, | |
| "gated": sum(len(r.gated) for r in reports), | |
| "collapsed": sum(len(r.collapsed) for r in reports), | |
| "low_information": sum(len(r.low_information) for r in reports), | |
| # classified, with fewer than SHORT_READ_BASES valid bases | |
| "short_read": sum(len(r.short_read) for r in reports), | |
| # species confidence at or above the threshold, species Unclassified for another reason | |
| "higher_rank": sum(len(r.higher_rank) for r in reports), | |
| "rounding": sum(len(r.rounding) for r in reports), | |
| "other_rule": sum(len(r.other_rule) for r in reports), | |
| # per rank: reads at or above the threshold without a label, by reason | |
| "withheld_by_rank": withheld_by_rank(reports), | |
| "low_information_limits": list(first.low_information_limits) if first is not None else None, | |
| # checked once, on the first chunk, which holds the reads the CLI checks | |
| "strand": list(first.strand) if first is not None and first.strand else None, | |
| } | |
| def withheld_by_rank(reports: Sequence[RunReport]) -> Dict[str, Dict[str, int]]: | |
| """Per rank and reason (``package_adapter.WITHHELD_REASONS``), the reads whose | |
| confidence reached the rank's threshold and that are Unclassified there, | |
| summed over the chunks of a job. Every such read is counted once per rank.""" | |
| out = {lvl: {k: 0 for k in package_adapter.WITHHELD_REASONS} for lvl in LEVEL_ORDER} | |
| for r in reports: | |
| for lvl, counts in r.withheld_by_rank.items(): | |
| for k, v in counts.items(): | |
| out[lvl][k] += v | |
| return out | |