A script to create a ranked embedding pickle with BioBert.

bert-embedding.py

from bert_serving.client import BertClient
from scipy import spatial
import pickle

bc = BertClient()

# load embeddings of the human subset of the UniProt entries
with open('/path/to/uniprot/protiens-/prots_human.tsv') as ref_file:
    ref_rows = ref_file.read().splitlines()[1:]
    print(ref_rows[:10])

ref_prots = [r.split('\t')[2] for r in ref_rows]
print(ref_prots[:10])

# get embeddings for ref prots
ref_embeddings = bc.encode(ref_prots)

print(len(ref_embeddings), len(ref_prots))

# get embedding for query
with open('/path/to/proteins/to/query-/bioc-train.csv') as query_file:
    queries = query_file.read().splitlines()[1:]

query_embeddings = bc.encode(queries)

# find and rank cosine distances
results = {}
for q_idx, query_embedding in enumerate(query_embeddings):
    print(q_idx, queries[q_idx])
    cos_distances = []
    for embed_idx, ref_embedding in enumerate(ref_embeddings):
        cos_distances.append((embed_idx, spatial.distance.cosine(query_embedding, ref_embedding)))

    # print(cos_distances[:10])

    results[queries[q_idx]] = [(ref_prots[x[0]], x[1]) for x in sorted(cos_distances, key=lambda t: t[1])[:10]]
    # print(results)

with open('/output/results.pkl', 'wb') as wfile:
    pickle.dump(results, wfile)