| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
|
|
| import diagindex as di |
| import dlmodel as dl |
| import mmsearch as ms |
| from inputs import prepare_inputs |
|
|
| BETA, GATE, BUDGET, SHARE = 0.7, 20, 420, 0.95 |
|
|
|
|
| def red_votes(index, lab_flat, lab_ptr, seqs, n_ec): |
| out = np.zeros((len(seqs), n_ec), dtype=np.float32) |
| best = np.zeros(len(seqs)) |
| for r, seq in enumerate(seqs): |
| cand, s = di.candidates(index, di.to_codes(seq)) |
| out[r] = di.vote(cand, s, lab_flat, lab_ptr, n_ec) |
| best[r] = s[0] if len(s) else 0 |
| return out, best |
|
|
|
|
| def predict_file(input_path, output_path): |
| torch.set_num_threads(2) |
| root = Path(__file__).parent |
| frame = prepare_inputs(pd.read_parquet(input_path)) |
| seqs = frame.sequence.tolist() |
| ec_list = json.loads((root / 'weights' / 'meta.json').read_text())['ec_list'] |
| n_ec = len(ec_list) |
|
|
| z = np.load(root / 'weights' / 'kmer_index.npz') |
| codes, lens = z['codes'], z['lens'] |
| lab_flat = z['lab_flat'].astype(np.int64) |
| lab_ptr = np.concatenate([[0], np.cumsum(z['lab_len'])]) |
|
|
| try: |
| hits = ms.search(root, codes, lens, seqs, threads=2, budget=BUDGET) |
| score, has = ms.votes(hits, lab_flat, lab_ptr, len(seqs), n_ec) |
| except Exception as error: |
| print('mmseqs failed, using the fallback for every input:', error, flush=True) |
| score, has = np.zeros((len(seqs), n_ec), dtype=np.float32), np.zeros(len(seqs), dtype=bool) |
|
|
| unsure = np.where(has & (score.max(1) < SHARE))[0] |
| nohit = np.where(~has)[0] |
| print('unsure inputs:', len(unsure), ' inputs without a search hit:', len(nohit), ' of', len(seqs), flush=True) |
|
|
| need = np.concatenate([unsure, nohit]).astype(np.int64) |
| if len(need): |
| index = di.build(codes, lens) |
| red, best = red_votes(index, lab_flat, lab_ptr, [seqs[i] for i in need], n_ec) |
| score[unsure] += red[:len(unsure)] |
| if len(nohit): |
| red_n, best_n = red[len(unsure):], best[len(unsure):] |
| weak = np.where(best_n < GATE)[0] |
| if len(weak): |
| P, cfg, bank, alpha = dl.load(root / 'weights') |
| emb, probs = dl.run(P, cfg, [seqs[nohit[i]] for i in weak], 'cpu') |
| cnn = alpha * dl.knn_votes(bank, emb) + (1 - alpha) * probs |
| red_n[weak] = BETA * red_n[weak] + (1 - BETA) * cnn |
| score[nohit] = red_n |
|
|
| preds = [ec_list[i] for i in score.argmax(axis=1)] |
| valid = set(json.loads((root / 'ec_vocabulary.json').read_text())) |
| assert set(preds) <= valid, 'prediction outside ec_vocabulary.json' |
| with Path(output_path).open('w') as output: |
| for key, ec in zip(frame.id, preds): |
| output.write(json.dumps({'id': key, 'prediction': [ec]}) + '\n') |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--input', required=True) |
| parser.add_argument('--output', required=True) |
| args = parser.parse_args() |
| predict_file(args.input, args.output) |
|
|
