File size: 4,535 Bytes
c56bdb6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | 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']), '--num-iterations', str(cfg.get('iters', 1)), '-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: # never fail because of the re-ranker: fall back to the vote
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) # identical inputs count once, first ID kept
seqs = uniq['sequence'].tolist()
order = sorted(range(len(seqs)), key=lambda i: len(seqs[i])) # short first: most inputs finish if time runs out
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()
|