| """Offline CPU inference: k-mer prefilter (5-mers + 4-mers) -> Smith-Waterman re-ranking -> top-20 weighted EC vote.""" |
| import argparse |
| import gc |
| import json |
| import os |
| import time |
| from collections import Counter, defaultdict |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import pyopal |
|
|
| from inputs import prepare_inputs |
| from kmer import KmerIndex |
|
|
| ROOT = Path(__file__).parent |
| T0 = time.time() |
| BUDGET = float(os.environ.get('TIME_BUDGET', 480)) |
| N5 = 1000 |
| N4 = 300 |
| TOP_K = 20 |
| EXPONENT = 16 |
| MAX_LEN_RATIO = 3 |
| THREADS = 2 |
|
|
|
|
| def log(msg): |
| print(f'[{time.time() - T0:6.1f}s] {msg}', flush=True) |
|
|
|
|
| def load_db(): |
| path = ROOT / 'weights/db.parquet' |
| if not path.exists(): |
| path = ROOT / 'db.parquet' |
| db = pd.read_parquet(path) |
| vocab = set(json.loads((ROOT / 'ec_vocabulary.json').read_text())) |
| ecs = [[e for e in row if e in vocab] for row in db.ecs] |
| keep = [i for i, e in enumerate(ecs) if e] |
| seqs = [db.sequence.iat[i] for i in keep] |
| ecs = [ecs[i] for i in keep] |
| fallback = Counter(e for row in ecs for e in row).most_common(1)[0][0] |
| return seqs, ecs, fallback |
|
|
|
|
| def vote(cand, scores, ecs, fallback): |
| if len(cand) == 0 or scores.max() <= 0: |
| return fallback |
| order = np.argsort(-scores)[:TOP_K] |
| best = scores[order[0]] |
| votes = defaultdict(float) |
| for i in order: |
| weight = (max(scores[i], 0) / best) ** EXPONENT |
| row = ecs[cand[i]] |
| for e in row: |
| votes[e] += weight / len(row) |
| return max(votes, key=votes.get) |
|
|
|
|
| def candidates(seqs, queries, k, n): |
| index = KmerIndex(seqs, k=k) |
| idx, _ = index.topk(queries, n) |
| del index |
| gc.collect() |
| log(f'{k}-mer candidates ready') |
| return idx |
|
|
|
|
| def predict_file(input_path, output_path): |
| frame = prepare_inputs(pd.read_parquet(input_path)) |
| queries = [s if isinstance(s, str) else '' for s in frame.sequence] |
| log(f'{len(queries)} unique inputs') |
|
|
| seqs, ecs, fallback = load_db() |
| c5 = candidates(seqs, queries, 5, N5) |
| |
| preds = [ecs[i[0]][0] if i[0] >= 0 else fallback for i in c5] |
| c4 = candidates(seqs, queries, 4, N4) if time.time() - T0 < BUDGET / 3 else None |
|
|
| aligner = pyopal.Aligner('BLOSUM62', gap_open=11, gap_extend=1) |
| database = pyopal.Database(seqs) |
| lengths = np.array([len(s) for s in seqs]) |
|
|
| def align(q): |
| if time.time() - T0 > BUDGET or not queries[q]: |
| return q, None |
| pool = c5[q] |
| if c4 is not None: |
| extra = c4[q][c4[q] >= 0] |
| extra = extra[lengths[extra] <= MAX_LEN_RATIO * len(queries[q])] |
| pool = np.concatenate([pool, extra]) |
| cand = [int(c) for c in dict.fromkeys(pool.tolist()) if c >= 0] |
| if not cand: |
| return q, None |
| scores = np.array([r.score for r in aligner.align(queries[q], database.extract(cand))], float) |
| return q, vote(cand, scores, ecs, fallback) |
|
|
| done = 0 |
| with ThreadPoolExecutor(THREADS) as pool: |
| for q, p in pool.map(align, range(len(queries))): |
| if p is not None: |
| preds[q] = p |
| done += 1 |
| log(f'aligned {done}/{len(queries)} queries') |
|
|
| with Path(output_path).open('w') as out: |
| for key, p in zip(frame.id, preds): |
| out.write(json.dumps({'id': key, 'prediction': [p]}) + '\n') |
| log(f'wrote {output_path}') |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--input', required=True) |
| parser.add_argument('--output', required=True) |
| args = parser.parse_args() |
| predict_file(args.input, args.output) |
|
|
