Make working embedding view

app.py

@@ -44,4 +44,4 @@ with tab1:
     )
 
 with tab2:
-    views.diffs(embeddings, corrector)
+    views.diffs(embeddings, corrector, encoder, tokenizer)

resources.py

@@ -1,10 +1,11 @@
 import streamlit as st
 import pandas as pd
+import torch
 import vec2text
 from transformers import AutoModel, AutoTokenizer
 from sklearn.decomposition import PCA
 from utils import file_cache
-
+from transformers import PreTrainedModel, PreTrainedTokenizer
 
 # Caching the vec2text corrector
 @st.cache_resource
@@ -35,4 +36,22 @@ def reduce_embeddings(embeddings):
 def load_model_and_tokenizer(device="cpu"):
     encoder = AutoModel.from_pretrained("sentence-transformers/gtr-t5-base").encoder.to(device)
     tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/gtr-t5-base")
-    return encoder, tokenizer
+    return encoder, tokenizer
+
+
+def get_gtr_embeddings(text_list,
+                       encoder: PreTrainedModel,
+                       tokenizer: PreTrainedTokenizer) -> torch.Tensor:
+
+    inputs = tokenizer(text_list,
+                       return_tensors="pt",
+                       max_length=128,
+                       truncation=True,
+                       padding="max_length",).to("cuda")
+
+    with torch.no_grad():
+        model_output = encoder(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
+        hidden_state = model_output.last_hidden_state
+        embeddings = vec2text.models.model_utils.mean_pool(hidden_state, inputs['attention_mask'])
+
+    return embeddings

views.py

@@ -5,16 +5,45 @@ from umap import UMAP
 import plotly.express as px
 import numpy as np
 from streamlit_plotly_events import plotly_events
-from resources import reduce_embeddings
 import utils
 import pandas as pd
 from scipy.spatial import distance
 
 dimensionality_reduction_model_name = "PCA"
 
-def diffs(embeddings: np.ndarray, corrector):
-    st.text(f"Embedding shape: {embeddings.shape}")
-    st.html('<a href="https://www.flaticon.com/free-icons/array" title="array icons">Array icons created by Voysla - Flaticon</a>')
+def diffs(embeddings: np.ndarray, corrector, encoder, tokenizer):
+    st.title('"A man is to king, what woman is to queen"')
+    st.markdown("A well known pehnomenon in semantic vectors is the way we can do vector operations like addition and subtraction to find spacial relations in the vector space.")
+    st.markdown(
+        'In word embedding models, we have found that the relationship between words can be captured mathematically, '
+        'such that "king" is to "man" as "queen" is to "woman," demonstrating that vector arithmetic can encode analogies and semantic relationships in high-dimensional space ([Mikolov et al., 2013](https://arxiv.org/abs/1301.3781)).'
+        )
+    st.markdown("This application lets you freely explore to which extent that property applies to embedding inversion models given the other factors of inaccuracy")
+
+    generated_sentence = ""
+
+
+    with st.form(key="foo") as form:
+        submit_button = st.form_submit_button("Synthesize")
+
+        sent1 = st.text_input("Sentence 1")
+        st.latex("-")
+        sent2 = st.text_input("Sentence 2")
+        st.latex("+")
+        sent3 = st.text_input("Sentence 3")
+        st.latex("=")
+
+        if submit_button:
+            generated_sentence = "HI"
+        
+        sent4 = st.text_input("Sentence 4", value=generated_sentence, disabled=True)
+
+        
+
+        if submit_button:
+            generated_sentence = "HI!"
+
+    # st.html('<a href="https://www.flaticon.com/free-icons/array" title="array icons">Array icons created by Voysla - Flaticon</a>')
 
 def plot(df: pd.DataFrame, embeddings: np.ndarray, vectors_2d, reducer, corrector):
     
@@ -88,7 +117,7 @@ def plot(df: pd.DataFrame, embeddings: np.ndarray, vectors_2d, reducer, correcto
 
         if inferred_embedding is not None and (closest_sentence_index != -1):
             couple = selected_sentence_embedding.squeeze(), inferred_embedding.squeeze()
-            st.markdown(f"### Inferred embedding distance:")
+            st.markdown("### Inferred embedding distance:")
             st.number_input("Euclidean", value=distance.euclidean(
                 *couple
             ), disabled=True)