<jupyter_start><jupyter_text>Setup<jupyter_code>import torch
from torch import einsum
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F
from clasp import CLASP, Transformer, tokenize, basic_rand_sampler, basic_aa_tokenizer, CLASPDataset
from functools import partial
from collections import OrderedDict
import requests
import matplotlib.pyplot as plt<jupyter_output><empty_output><jupyter_text>Instantiate model<jupyter_code>text_enc = Transformer(
    num_tokens = 49408,
    dim = 512,
    depth = 6,
    seq_len = 1024
)

bioseq_enc = Transformer(
    num_tokens = 23,
    dim = 512,
    depth = 6,
    seq_len = 512,
    sparse_attn = False
)

model = CLASP(
    text_encoder = text_enc,
    bioseq_encoder = bioseq_enc
)
path_weights = "xy/latest.pt"
ckpt = torch.load(path_weights, map_location="cpu")
# this hack seems to be needed, because the model in the ipynb has different weight names?
new_ckpt = OrderedDict()
for ((kc, vc), ko) in zip(ckpt.items(), model.state_dict().keys()):
    new_ckpt[ko] = vc
model.load_state_dict(new_ckpt)
model.eval();
# if you wanto to run it on GPU
#device = torch.device("cuda:0")
#model.to(device)<jupyter_output><empty_output><jupyter_text>Tokenizer setup<jupyter_code>text_sampler = partial(basic_rand_sampler, sample_len=1024)
text_tok = partial(tokenize, context_length=1024, return_mask=True)

bioseq_sampler = partial(basic_rand_sampler, sample_len=512)
bioseq_tok = partial(basic_aa_tokenizer, context_length=512, return_mask=True)<jupyter_output><empty_output><jupyter_text>Query example  Bioseq Proteins related query "blood":1. https://www.uniprot.org/uniprot/B2RUZ41. https://www.uniprot.org/uniprot/A0A430QLP7Proteins related to query "brain":1. https://www.uniprot.org/uniprot/Q9VGE81. https://www.uniprot.org/uniprot/Q8C8N3Proteins related to query "t-cell":1. https://www.uniprot.org/uniprot/P209631. https://www.uniprot.org/uniprot/P60568Proteins related to query "hormones":1. https://www.uniprot.org/uniprot/P012151. https://www.uniprot.org/uniprot/O97385Proteins related to query "virus":1. https://www.uniprot.org/uniprot/P042981. https://www.uniprot.org/uniprot/Q8TDQ0<jupyter_code>prot_urls = ["https://www.uniprot.org/uniprot/B2RUZ4.fasta",
             "https://www.uniprot.org/uniprot/A0A430QLP7.fasta",
             "https://www.uniprot.org/uniprot/Q9VGE8.fasta",
             "https://www.uniprot.org/uniprot/Q8C8N3.fasta",
             "https://www.uniprot.org/uniprot/P20963.fasta",
             "https://www.uniprot.org/uniprot/P60568.fasta",
             "https://www.uniprot.org/uniprot/O97385.fasta",
             "https://www.uniprot.org/uniprot/P04298.fasta",
             "https://www.uniprot.org/uniprot/Q8TDQ0.fasta"]
bioseq_orig =["".join(requests.get(url).text.split("\n")[1:]) for url in prot_urls]
bioseq_orig
[len(s) for s in bioseq_orig]
bioseq_ = [bioseq_tok(bioseq_sampler(s)) for s in bioseq_orig]
bioseq, bioseq_mask = list(zip(*bioseq_))
bioseq      = torch.cat([s.unsqueeze(0) for s in bioseq])
bioseq_mask = torch.cat([s.unsqueeze(0) for s in bioseq_mask])
bioseq, bioseq_mask
bioseq.shape, bioseq_mask.shape
bioseq_out = model.bioseq_encoder(bioseq, mask=bioseq_mask)
bioseq_out.shape<jupyter_output><empty_output><jupyter_text>Text Change here the text query:<jupyter_code>#text_orig = "protein that is present in the blood"
text_orig = "blood cells blood stream nutrition"
#text_orig = "brain neurology nerve cells"
#text_orig = "proteins that play a role in the blood transport"
#text_orig = "hormone regulation and similar functions"
#text_orig = "immune system and antigen processing"
#text_orig = "viral infection and replication"
text, text_mask = text_tok(text_orig)
text, text_mask
text_out = model.text_encoder(text, mask=text_mask)
text_out.shape<jupyter_output><empty_output><jupyter_text>Similarity<jupyter_code>text_out_norm = F.normalize(text_out, p=2, dim =-1)
bioseq_out_norm = F.normalize(bioseq_out, p=2, dim =-1)
sim = einsum('n d, n d -> n', text_out_norm, bioseq_out_norm)
sim.tolist()
plt.imshow(sim.data.unsqueeze(0).numpy(), cmap="turbo");
plt.imshow(sim.data.unsqueeze(0).numpy(), cmap="turbo", vmin=-1., vmax=1.);<jupyter_output><empty_output><jupyter_text>Citric acid cycle The [citric acid cycle](https://en.wikipedia.org/wiki/Citric_acid_cycle) is a very import chain of reactions in cells and we use the sequences of some of the [involved enzymes](https://en.wikipedia.org/wiki/Citric_acid_cycleSteps) to look for interesting relationships in the high-dimensional embedding space.1. Citrate synthase, https://www.uniprot.org/uniprot/O753901. Cytoplasmic aconitate hydratase, https://www.uniprot.org/uniprot/P213991. Isocitrate dehydrogenase NADP cytoplasmic, https://www.uniprot.org/uniprot/O758741. Fumarate hydratase mitochondrial, https://www.uniprot.org/uniprot/P079541. Malate dehydrogenase cytoplasmic, https://www.uniprot.org/uniprot/P40925<jupyter_code>cca_urls = ["https://www.uniprot.org/uniprot/O75390.fasta",
            "https://www.uniprot.org/uniprot/P21399.fasta",
            "https://www.uniprot.org/uniprot/O75874.fasta",
            "https://www.uniprot.org/uniprot/P07954.fasta",
            "https://www.uniprot.org/uniprot/P40925.fasta"]
# we also add the proteins from above to see if the cca proteins are clustered together
cca_seq =["".join(requests.get(url).text.split("\n")[1:]) for url in cca_urls+prot_urls]
cca_seq
[len(s) for s in cca_seq]
bioseq_ = [bioseq_tok(bioseq_sampler(s)) for s in cca_seq]
bioseq, bioseq_mask = list(zip(*bioseq_))
bioseq      = torch.cat([s.unsqueeze(0) for s in bioseq])
bioseq_mask = torch.cat([s.unsqueeze(0) for s in bioseq_mask])
bioseq, bioseq_mask
bioseq.shape, bioseq_mask.shape
bioseq_out = model.bioseq_encoder(bioseq, mask=bioseq_mask)
bioseq_out.shape
u, s, v = torch.svd(bioseq_out)
#u, s, v
s
u.shape, s.shape, v.shape
bioseq_out.shape, v.t()[0].unsqueeze(0).shape
# project on the top-2 eigenvectors
x_coords = einsum("id, jd -> ij", bioseq_out, v.t()[0].unsqueeze(0)).data
y_coords = einsum("id, jd -> ij", bioseq_out, v.t()[1].unsqueeze(0)).data
plt.scatter(x_coords, y_coords);
for i, txt in enumerate(x_coords):
    plt.annotate(i, (x_coords[i]+0.025, y_coords[i]-0.2))<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup