Implement genus classification

app.py

@@ -1,3 +1,5 @@
+from io import BytesIO
+import os
 import re
 import PIL.Image
 import pandas as pd
@@ -12,14 +14,18 @@ import torch
 from torch import nn
 from transformers import BertConfig, BertForMaskedLM, PreTrainedTokenizerFast
 from huggingface_hub import PyTorchModelHubMixin
+from pinecone import Pinecone
 
-from config import DEFAULT_INPUTS, MODELS, DATASETS
+from config import DEFAULT_INPUTS, MODELS, DATASETS, ID_TO_GENUS_MAP
 
 # We need this for the eco layers because they are too big
 PIL.Image.MAX_IMAGE_PIXELS = None
 
 torch.set_grad_enabled(False)
 
+# Configure pinecone
+pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
+pc_index = pc.Index("amazon")
 
 # Load models
 class DNASeqClassifier(nn.Module, PyTorchModelHubMixin):
@@ -60,10 +66,8 @@ def set_default_inputs():
             DEFAULT_INPUTS["longitude"])
 
 
-def preprocess(dna_sequence: str, latitude: str, longitude: str): 
-    """
-    Prepares app input for downsteram tasks
-    """
+def preprocess(dna_sequence: str, latitude: float, longitude: float): 
+    """Prepares app input for downsteram tasks"""
 
     # Preprocess the DNA sequence turning it into an embedding
     dna_seq_preprocessed: str = re.sub(r"[^ACGT]", "N", dna_sequence)
@@ -80,58 +84,65 @@ def preprocess(dna_sequence: str, latitude: str, longitude: str):
     # Preprocess the location data
     coords = (float(latitude), float(longitude))
 
-    return dna_embedding, coords
-    # ecolayer_data = ecolayers_ds  # TODO something something...
+    return dna_embedding, coords[0], coords[1]
 
-    # # format lat and lon into coords
-    # coords = (inp_lat, inp_lng)
-    # # Grab rasters from the tifs
-    # ecoLayers = load_dataset("LofiAmazon/Global-Ecolayers")
-    # temp = pd.DataFrame([coords, embed], columns = ['coord', 'embeddings'])
-    # data = pd.merge(temp, ecoLayers, on='coord', how='left')
 
-    # return data
-
-# def predict_genus():
-#     data = preprocess()
-#     out = infer.infer_dna(data)
-    
-#     results = []
-
-#     genuses = infer.infer()
-        
-#     results.append({
-#         "sequence": dna_df['nucraw'],
-#         # "predictions": pd.concat([dna_genuses, envdna_genuses], axis=0)
-#         'predictions': genuses})
-    
-#     return results
+def tokenize(dna_sequence: str) -> dict[str, torch.Tensor]:
+    dna_seq_preprocessed: str = re.sub(r"[^ACGT]", "N", dna_sequence)
+    dna_seq_preprocessed: str = re.sub(r"N+$", "", dna_sequence)
+    dna_seq_preprocessed = dna_seq_preprocessed[:660]
+    dna_seq_preprocessed = " ".join([
+        dna_seq_preprocessed[i:i+4] for i in range(0, len(dna_seq_preprocessed), 4)
+    ])
 
-# def tsne_DNA(data, genuses):
-#     data["embeddings"] = data["embeddings"].apply(lambda x: np.array(list(map(float, x[1:-1].split()))))
+    return tokenizer(dna_seq_preprocessed, return_tensors="pt")
 
-#     # Pick genuses with most samples
-#     top_k = 5
-#     genus_counts = df["genus"].value_counts()
-#     top_genuses = genus_counts.head(top_k).index
-#     df = df[df["genus"].isin(top_genuses)]
 
-#     # Create a t-SNE plot of the embeddings
-#     n_genus = len(df["genus"].unique())
-#     tsne = TSNE(n_components=2, perplexity=30, learning_rate=200, n_iter=1000, random_state=0)
 
-#     X = np.stack(df["embeddings"].tolist())
-#     y = df["genus"].tolist()
+def get_embedding(dna_sequence: str) -> torch.Tensor:
+    dna_embedding: torch.Tensor = embeddings_model(
+        **tokenize(dna_sequence)
+    ).hidden_states[-1].mean(1).squeeze()
 
-#     X_tsne = tsne.fit_transform(X)
+    return dna_embedding
 
-#     label_encoder = LabelEncoder()
-#     y_encoded = label_encoder.fit_transform(y)
 
-#     plot = plt.figure(figsize=(6, 5))
-#     scatter = plt.scatter(X_tsne[:, 0], X_tsne[:, 1], c=y_encoded, cmap="viridis", alpha=0.7)
+def predict_genus(method: str, dna_sequence: str, latitude: str, longitude: str):
+    coords = (float(latitude), float(longitude))
 
-#     return plot
+    if method == "cosine":
+        embedding = get_embedding(dna_sequence)
+        result = pc_index.query(
+            namespace="all",
+            vector=embedding.tolist(),
+            top_k=100,
+            include_metadata=True,
+        )
+        top_k = [m["metadata"]["genus"] for m in result["matches"]]
+    
+        top_k = pd.Series(top_k).value_counts()
+        top_k = top_k / top_k.sum()
+    
+    if method == "fine_tuned_model":
+        bert_inputs = tokenize(dna_sequence)
+        logits = classification_model(bert_inputs, torch.zeros(1, 7))
+        probs = torch.softmax(logits, dim=1).squeeze()
+        top_k = torch.topk(probs, 10)
+        top_k = pd.Series(
+            top_k.values.detach().numpy(),
+            index=[ID_TO_GENUS_MAP[i] for i in top_k.indices.detach().numpy()]
+        )
+        # top_k = pd.Series(top_k.values.detach().numpy(), index=top_k.indices.detach().numpy())
+
+    fig, ax = plt.subplots()
+    ax.bar(top_k.index.astype(str), top_k.values)
+    ax.set_title("Genus Prediction")
+    ax.set_xlabel("Genus")
+    ax.set_ylabel("Probability")
+    ax.set_xticklabels(top_k.index.astype(str), rotation=90) 
+    fig.canvas.draw()
+
+    return PIL.Image.frombytes("RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb())
 
 
 with gr.Blocks() as demo:
@@ -156,20 +167,34 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         btn_run = gr.Button("Predict")
-        btn_run.click(fn=preprocess, inputs=[inp_dna, inp_lat, inp_lng])
+        btn_run.click(
+            fn=preprocess,
+            inputs=[inp_dna, inp_lat, inp_lng],
+        )
 
         btn_defaults = gr.Button("I'm feeling lucky")
         btn_defaults.click(fn=set_default_inputs, outputs=[inp_dna, inp_lat, inp_lng])
 
-
     with gr.Tab("Genus Prediction"):
-        with gr.Row():
-            gr.Markdown("Make plot or table for Top 5 species")
+        gr.Interface(
+            fn=predict_genus,
+            inputs=[
+                gr.Dropdown(choices=["cosine", "fine_tuned_model"], value="fine_tuned_model"),
+                inp_dna,
+                inp_lat,
+                inp_lng,
+            ],
+            outputs=["image"],
+        )
+
+        # with gr.Row():
+
+        #     gr.Markdown("Make plot or table for Top 5 species")
+
+        # with gr.Row():
+        #     genus_out = gr.Dataframe(headers=["DNA Only Pred Genus", "DNA Only Prob", "DNA & Env Pred Genus", "DNA & Env Prob"])
+        #     # btn_run.click(fn=predict_genus, inputs=[inp_dna, inp_lat, inp_lng], outputs=genus_out)
 
-        with gr.Row():
-            genus_out = gr.Dataframe(headers=["DNA Only Pred Genus", "DNA Only Prob", "DNA & Env Pred Genus", "DNA & Env Prob"])
-            # btn_run.click(fn=predict_genus, inputs=[inp_dna, inp_lat, inp_lng], outputs=genus_out)
-    
     with gr.Tab('DNA Embedding Space Visualizer'):
         gr.Markdown("If the highest genus probability is very low for your DNA sequence, we can still examine the DNA embedding of the sequence in relation to known samples for clues.")
 

requirements.txt

@@ -1,5 +1,6 @@
 huggingface-hub==0.23.2
 pandas==2.2.2
+pinecone_client==4.1.0
 rasterio==1.3.10
 torch==2.3.0
 tqdm==4.66.4