| import argparse |
| import json |
| import subprocess |
| import sys |
| import tempfile |
| import time |
| from pathlib import Path |
|
|
| import pandas as pd |
|
|
| HERE = Path(__file__).resolve().parent |
| sys.path.insert(0, str(HERE)) |
| import knn |
| import rerank |
|
|
|
|
| def search_chunk(rows, seqs, cfg): |
| with tempfile.TemporaryDirectory(dir='/tmp') as td: |
| td = Path(td) |
| with open(td / 'q.fa', 'w') as f: |
| for i in rows: |
| f.write('>q%d\n%s\n' % (i, seqs[i])) |
| out = td / 'hits.tsv' |
| cmd = [str(HERE / 'bin' / 'mmseqs'), 'easy-search', str(td / 'q.fa'), str(HERE / 'db' / 'trainDB'), str(out), |
| str(td / 'work'), '-s', str(cfg['sens']), '--max-seqs', str(cfg['max_seqs']), '-e', '10', |
| '--split-memory-limit', cfg['split_mem'], '--threads', str(cfg['threads']), |
| '--format-output', ','.join(rerank.HIT_FIELDS), '-v', '1'] |
| r = subprocess.run(cmd, capture_output=True, text=True) |
| if r.returncode != 0: |
| sys.stderr.write(r.stdout[-2000:] + '\n' + r.stderr[-2000:] + '\n') |
| sys.exit('mmseqs failed with code %d' % r.returncode) |
| if out.exists() and out.stat().st_size > 0: |
| h = pd.read_csv(out, sep='\t', names=rerank.HIT_FIELDS) |
| if h['pident'].max() > 1.0: |
| h['pident'] = h['pident'] / 100 |
| return h |
| return pd.DataFrame(columns=rerank.HIT_FIELDS) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--input', required=True) |
| ap.add_argument('--output', required=True) |
| args = ap.parse_args() |
| t0 = time.perf_counter() |
| cfg = json.loads((HERE / 'config.json').read_text()) |
| ec_map = json.loads((HERE / 'db' / 'ec_map.json').read_text()) |
| fallback = json.loads((HERE / 'db' / 'fallback.json').read_text()) |
| freq = json.loads((HERE / 'db' / 'freq.json').read_text()) |
| l3_tot = rerank.l3_totals(freq) |
|
|
| model = None |
| if cfg.get('use_rerank'): |
| try: |
| import joblib |
| model = joblib.load(HERE / 'model' / 'rr.joblib') |
| except Exception as e: |
| sys.stderr.write('re-ranker unavailable, using the vote only: %r\n' % (e,)) |
|
|
| df = pd.read_parquet(args.input, columns=['id', 'sequence']) |
| uniq = df.drop_duplicates('sequence', keep='first').reset_index(drop=True) |
| seqs = uniq['sequence'].tolist() |
| order = sorted(range(len(seqs)), key=lambda i: len(seqs[i])) |
| chunks = [order[i:i + cfg['chunk']] for i in range(0, len(order), cfg['chunk'])] |
|
|
| hl, rr, times, skipped = {}, {}, [], 0 |
| for ci, rows in enumerate(chunks): |
| elapsed = time.perf_counter() - t0 |
| if times and elapsed + 1.3 * max(times) > cfg['time_budget']: |
| skipped = sum(len(c) for c in chunks[ci:]) |
| sys.stderr.write('TIME GUARD: stopping before chunk %d; %d inputs get the fallback EC\n' % (ci, skipped)) |
| break |
| tc = time.perf_counter() |
| hits = search_chunk(rows, seqs, cfg) |
| hl.update(knn.hits_from_df(hits[['query', 'target', 'bits']], cfg['kmax'])) |
| if model is not None: |
| try: |
| cand = rerank.build_candidates(hits, ec_map, freq, cfg['kmax'], l3_tot) |
| rr.update(rerank.rerank_select(cand, model, cfg['rel'], cfg['max_ec'])) |
| except Exception as e: |
| sys.stderr.write('re-ranker failed, using the vote for the rest: %r\n' % (e,)) |
| model = None |
| times.append(time.perf_counter() - tc) |
| sys.stderr.write('chunk %d/%d done in %.0fs (elapsed %.0fs)\n' % (ci + 1, len(chunks), times[-1], time.perf_counter() - t0)) |
|
|
| changed = 0 |
| with open(args.output, 'w') as f: |
| for i, rid in enumerate(uniq['id']): |
| q = 'q%d' % i |
| vote = knn.select_ecs(hl.get(q, []), ec_map, fallback, cfg['K'], cfg['pw'], cfg['ratio'], cfg['max_ec'], freq) |
| ecs = rr.get(q) or vote |
| changed += int(ecs[0] != vote[0]) |
| f.write(json.dumps({'id': rid, 'prediction': ecs}) + '\n') |
| sys.stderr.write('re-ranker used: %s | first EC differs from the vote for %d inputs\n' % (bool(rr), changed)) |
| sys.stderr.write('total %.0fs for %d unique inputs (%d without search)\n' % (time.perf_counter() - t0, len(seqs), skipped)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
