added tsne graph, chanaged default coords

.gitignore

@@ -1,4 +1,5 @@
 .venv
 flagged
 *.tif
-*.tiff
+*.tiff
+.env

app.py

@@ -60,7 +60,8 @@ embeddings_model.eval()
 classification_model.eval()
 
 # Load datasets
-amazon_ds = load_dataset(DATASETS["amazon"])
+amazon_ds = load_dataset(DATASETS["amazon"])['train'].to_pandas()
+amazon_ds = amazon_ds[amazon_ds["genus"].notna()]
 
 def set_default_inputs():
     return (DEFAULT_INPUTS["dna_sequence"],
@@ -148,6 +149,22 @@ def predict_genus(method: str, dna_sequence: str, latitude: str, longitude: str)
             index=[ID_TO_GENUS_MAP[i] for i in top_k.indices.detach().numpy()]
         )
 
+    # fig, ax = plt.subplots()
+    # ax.bar(top_k.index.astype(str), top_k.values)
+    # ax.set_ylim(0, 1)
+    # ax.set_title("Genus Prediction")
+    # ax.set_xlabel("Genus")
+    # ax.set_ylabel("Probability")
+    # ax.set_xticks(range(len(top_k)))
+    # ax.set_xticklabels(top_k.index.astype(str), rotation=90)
+    # fig.subplots_adjust(bottom=0.3)
+    # fig.canvas.draw()
+
+    # return PIL.Image.frombytes("RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb())
+    return top_k
+
+def genus_hist(method: str, dna_sequence: str, latitude: str, longitude: str):
+    top_k = predict_genus(method, dna_sequence, latitude, longitude)
     fig, ax = plt.subplots()
     ax.bar(top_k.index.astype(str), top_k.values)
     ax.set_ylim(0, 1)
@@ -162,12 +179,12 @@ def predict_genus(method: str, dna_sequence: str, latitude: str, longitude: str)
     return PIL.Image.frombytes("RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb())
 
 
-def cluster_dna(top_k: float):
-    df = amazon_ds["train"].to_pandas()
-    df = df[df["genus"].notna()]
-    top_k = int(top_k)
+def cluster_dna(k: float):
+    df = amazon_ds
+    # df = df[df["genus"].notna()]
+    k = int(k)
     genus_counts = df["genus"].value_counts()
-    top_genuses = genus_counts.head(top_k).index
+    top_genuses = genus_counts.head(k).index
     df = df[df["genus"].isin(top_genuses)]
     tsne = TSNE(
         n_components=2, perplexity=30, learning_rate=200,
@@ -180,16 +197,59 @@ def cluster_dna(top_k: float):
 
     label_encoder = LabelEncoder()
     y_encoded = label_encoder.fit_transform(y)
+    classes = list(label_encoder.inverse_transform(range(len(df['genus'].unique()))))
 
     fig, ax = plt.subplots()
-    ax.scatter(X_tsne[:, 0], X_tsne[:, 1], c=y_encoded, cmap="viridis", alpha=0.7)
-    ax.set_title(f"DNA Embedding Space (of {str(top_k)} most common genera)")
+    plot = ax.scatter(X_tsne[:, 0], X_tsne[:, 1], c=y_encoded, cmap="tab20", alpha=0.7)
+    handles, _ = plot.legend_elements(prop='colors')
+    ax.legend(handles, classes)
+    ax.set_title(f"DNA Embedding Space (of {str(k)} most common genera)")
     # Reduce unnecessary whitespace
     ax.set_xlim(X_tsne[:, 0].min() - 0.1, X_tsne[:, 0].max() + 0.1)
     fig.canvas.draw()
 
     return PIL.Image.frombytes("RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb())
 
+def cluster_dna2(k: float, method: str, dna_sequence: str, latitude: str, longitude: str):
+    top_genuses = predict_genus(method, dna_sequence, latitude, longitude)
+    embed = get_embedding(dna_sequence).tolist()
+    # df = amazon_ds["train"].to_pandas()
+    df = amazon_ds
+    # df = df[df["genus"].notna()]
+    k = int(k)
+    # genus_counts = df["genus"].value_counts()
+    top_genuses = top_genuses.head(k).index
+    df = df[df["genus"].isin(top_genuses)]
+    tsne = TSNE(
+        n_components=2, perplexity=30, learning_rate=200,
+        n_iter=1000, random_state=0,
+    )
+    X = np.vstack([df['embeddings'].tolist(), embed])
+    # X = np.stack(df["embeddings"].tolist())
+    y = df["genus"].tolist()
+
+    X_tsne = tsne.fit_transform(X)
+    tsne_embed_space = X_tsne[:-1]
+    tsne_single = X_tsne[-1]
+
+    label_encoder = LabelEncoder()
+    y_encoded = label_encoder.fit_transform(y)
+    classes = list(label_encoder.inverse_transform(range(len(df['genus'].unique()))))
+
+    fig, ax = plt.subplots()
+    plot = ax.scatter(tsne_embed_space[:, 0], tsne_embed_space[:, 1], c=y_encoded, cmap="tab20", alpha=0.7)
+    ax.scatter(tsne_single[0], tsne_single[1], color='red', edgecolor='black')
+    handles, _ = plot.legend_elements(prop='colors')
+    ax.legend(handles, classes)
+    # ax.legend(loc='best')
+    ax.text(tsne_single[0], tsne_single[1], 'Your DNA Seq', fontsize=10, color='black')
+    ax.set_title(f"DNA Embedding Space Around Your DNA's Embedding")
+    # Reduce unnecessary whitespace
+    ax.set_xlim(X_tsne[:, 0].min() + 0.1, X_tsne[:, 0].max() + 0.1)
+    fig.canvas.draw()
+
+    return PIL.Image.frombytes("RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb())
+
 with gr.Blocks() as demo:
     # Header section
     gr.Markdown(("""
@@ -209,9 +269,9 @@ with gr.Blocks() as demo:
 
         with gr.Column():
             with gr.Row():
-                inp_lat = gr.Textbox(label="Latitude", placeholder="e.g. -3.009083")    
+                inp_lat = gr.Textbox(label="Latitude", placeholder="e.g. 2.009083")    
             with gr.Row():
-                inp_lng = gr.Textbox(label="Longitude", placeholder="e.g. -58.68281")
+                inp_lng = gr.Textbox(label="Longitude", placeholder="e.g. -41.68281")
 
     with gr.Row():
         btn_defaults = gr.Button("I'm feeling lucky")
@@ -224,13 +284,12 @@ with gr.Blocks() as demo:
         A demo of predicting the genus of a DNA sequence using multiple
         approaches (method dropdown):
 
-        - **fine_tuned_model**: using our
-          `LofiAmazon/BarcodeBERT-Finetuned-Amazon` which predicts the genus
+        - **fine_tuned_model**: uses our
+          `LofiAmazon/BarcodeBERT-Finetuned-Amazon` model which predicts the genus
           based on the DNA sequence and environmental data.
         - **cosine**: computes a cosine similarity between the DNA sequence
           embedding generated by our model and the embeddings of known samples
-          that we precomputed and stored in a Pinecone index. Thie method
-          DOES NOT examine ecological layer data.
+          that we precomputed and stored. This method DOES NOT use ecological layer data.
         """)
 
         with gr.Row():
@@ -243,34 +302,66 @@ with gr.Blocks() as demo:
                 genus_output = gr.Image()
 
             predict_button.click(
-                fn=predict_genus,
+                fn=genus_hist,
                 inputs=[method_dropdown, inp_dna, inp_lat, inp_lng],
                 outputs=genus_output
             )
 
     with gr.Tab("DNA Embedding Space Visualizer"):
         gr.Markdown("""
-        ## DNA Embedding Space Visualizer
-                    
-        We show a 2D t-SNE plot of the DNA embeddings of the five most common
-        genera in our dataset. This shows that the DNA Transformer model is
-        learning to cluster similar DNA sequences together.
-        """)
+            ## DNA Embedding Space Visualizer
+                        
+            Use this tool to visualize how our DNA Transformer model
+            learns to cluster similar DNA sequences together.
+            """)
+
+        # with gr.Row():
+        #     with gr.Column():
+        #         top_k_slider = gr.Slider(
+        #             minimum=1, maximum=10, step=1, value=5,
+        #             label="Choose **k**, the number of top genera to visualize",
+        #         )
+        #         visualize_button = gr.Button("Visualize Embedding Space")
+        #     with gr.Column():
+        #         visualize_output = gr.Image()
+
+        #     visualize_button.click(
+        #         fn=cluster_dna,
+        #         inputs=top_k_slider,
+        #         outputs=visualize_output
+        #     )
 
+        with gr.Row():
+            top_k_slider = gr.Slider(
+                minimum=1, maximum=10, step=1, value=5,
+                label="Choose **k**, the number of top genera to visualize",
+            )
+            visualize_button = gr.Button("Visualize Embedding Space")
         with gr.Row():
             with gr.Column():
-                top_k_slider = gr.Slider(
-                    minimum=1, maximum=10, step=1, value=5,
-                    label="Number of top genera to visualize",
-                )
-                visualize_button = gr.Button("Visualize Embedding Space")
-            with gr.Column():
+                gr.Markdown("""
+                    t-SNE plot of the DNA embedding spaces of the **k** most common
+                    genera in our dataset.
+                    """)
                 visualize_output = gr.Image()
 
-            visualize_button.click(
-                fn=cluster_dna,
-                inputs=top_k_slider,
-                outputs=visualize_output
+                visualize_button.click(
+                    fn=cluster_dna,
+                    inputs=top_k_slider,
+                    outputs=visualize_output
             )
+            with gr.Column():
+                gr.Markdown("""                            
+                    t-SNE plot of the DNA embedding spaces of the **k** most likely 
+                    genera for the DNA sequence you provided.
+                    """)
+                visualize_output2 = gr.Image()
+                
+                visualize_button.click(
+                    fn=cluster_dna2,
+                    inputs=[top_k_slider, method_dropdown, inp_dna, inp_lat, inp_lng],
+                    outputs=visualize_output2
+            )
+
 
 demo.launch()

default_inputs.json

@@ -1,5 +1,5 @@
 {
     "dna_sequence": "AACAATGTATTTGATTTTCGCCCTTGTGAATTTATTCGCTGGCGGAACAATGGCATTGTTGATTCGTTTGGAGTTGTTCCAACCTGGCTTGCAATTTTTAAGACCTGAGTTTTTTAATCAGTTAACAACTATGCACGGCCTTATAATGGTTTTCGGTGCAATTATGCCGGCCTTTGTGGGTTTTGCTAACTTGATGATTCCTTTGCAAATTGGTGCCTCTGATATGGCGTTTGCAAGAATGAACAATTTTAGTTTCTGGATTATGCCTGTTGCAGGGATGTTATTATTTGGCTCATTTTTGGCTCCTGGTGGCGCTACTGCAGCTGGTTGGACTTTGTATGCTCCTTTGTCGGTCCAAATGGGGCCTGGTATGGACATGACTATTTTTGCTGTTCACTTGATGGGTGCTTCATCCATTATGGGATCCATTAATATCATTGTGACAATTCTGAATATGCGTGCTCCTGGACTGTCTTTGATGAAGATGCCAATGTTCTGTTGGACATGGTTGATTACTGCATATTTGTTAATTGCGGTTATGCCTGTTTTAGCTGGTGCTATCACTATGGTTCTAACAGACCGTCACTTTGGAACAAGCTTTTTTGCAGCTGCTGGCGGTGGAGACCCTGTAATGTATCAACATATCTTC",
-    "latitude": "-3.009083",
-    "longitude": "-58.68281"
+    "latitude": "2.009083",
+    "longitude": "-41.68281"
 }