File size: 3,035 Bytes
d789132 017b9f5 d789132 9401a93 d789132 0c2054a d789132 a774c5a 0c2054a a774c5a 0c2054a a774c5a 0c2054a a774c5a 017b9f5 d789132 017b9f5 0c2054a 017b9f5 0c2054a 017b9f5 0c2054a 9401a93 a774c5a 9401a93 d789132 | 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 | 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)
|