add wiki model

app.py

@@ -3,9 +3,27 @@ import jax.numpy as jnp
 from transformers import AutoTokenizer
 from transformers.models.t5.modeling_flax_t5 import shift_tokens_right
 from t5_vae_flax_alt.src.t5_vae import FlaxT5VaeForAutoencoding
+import info
 
 
-st.title('T5-VAE')
+st.set_page_config(
+    page_title="T5-VAE",
+    page_icon="🙂😐🙁",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+
+st.title('T5-VAE 🙂😐🙁')
+
+st.text('''
+This is a variational autoencoder trained on text.
+
+It allows interpolating on text at a high level, try it out!
+
+See how it works [here](http://fras.uk/ml/large%20prior-free%20models/transformer-vae/2020/08/13/Transformers-as-Variational-Autoencoders.html).
+''')
+
 st.text('''
 Try interpolating between lines of Python code using this T5-VAE.
 ''')
@@ -79,11 +97,13 @@ def slerp(ratio, t1, t2):
     return res
 
 
-def decode(ratio, txt_1, txt_2):
+def decode(cnt, ratio, txt_1, txt_2):
     if not txt_1 or not txt_2:
         return ''
+    cnt.write('Getting latents...')
     lt_1, lt_2 = get_latent(txt_1), get_latent(txt_2)
     lt_new = slerp(ratio, lt_1, lt_2)
+    cnt.write('Decoding latent...')
     tkns = tokens_from_latent(lt_new)
     return tokenizer.decode(tkns.sequences[0], skip_special_tokens=True)
 
@@ -93,6 +113,62 @@ in_2 = st.text_input("Another line of Python code.", "x = a + 10 * 2")
 r = st.slider('Interpolation Ratio',  min_value=0.0, max_value=1.0, value=0.5)
 container = st.empty()
 container.write('Loading...')
+out = decode(container, r, in_1, in_2)
+container.empty()
+st.write(out)
+
+
+st.text('''
+Try interpolating between sentences from wikipedia using this T5-VAE.
+''')
+
+
+@st.cache(allow_output_mutation=True)
+def get_wiki_model():
+    tokenizer = AutoTokenizer.from_pretrained("t5-base")
+    model = FlaxT5VaeForAutoencoding.from_pretrained("flax-community/t5-vae-wiki")
+    assert model.params['t5']['shared']['embedding'].shape[0] == len(tokenizer), "T5 Tokenizer doesn't match T5Vae embedding size."
+    return model, tokenizer
+
+
+model, tokenizer = get_wiki_model()
+
+
+in_1 = st.text_input("A sentence.", "Children are looking for the water to be clear.")
+in_2 = st.text_input("Another sentence.", "There are two people playing soccer.")
+r = st.slider('Interpolation Ratio',  min_value=0.0, max_value=1.0, value=0.5)
+container = st.empty()
+container.write('Loading...')
 out = decode(r, in_1, in_2)
 container.empty()
 st.write(out)
+
+
+st.text('''
+Try arithmetic in latent space.
+''')
+
+
+def arithmetic(cnt, txt_a, txt_b, txt_c):
+    if not txt_a or not txt_b or not txt_c:
+        return ''
+    cnt.write('getting latents...')
+    lt_a, lt_b, lt_c = get_latent(txt_a), get_latent(txt_b), get_latent(txt_c)
+    lt_d = lt_c + (lt_b - lt_a)
+    cnt.write('decoding C + (B - A)...')
+    tkns = tokens_from_latent(lt_d)
+    return tokenizer.decode(tkns.sequences[0], skip_special_tokens=True)
+
+
+in_a = st.text_input("A", "Children are looking for the water to be clear.")
+in_b = st.text_input("B", "There are two people playing soccer.")
+in_c = st.text_input("C", "Children are looking for the water to be clear.")
+
+st.text('''
+A is to B as C is to...
+''')
+container = st.empty()
+container.write('Loading...')
+out = arithmetic(container, in_a, in_b, in_c)
+container.empty()
+st.write(out)

info.py

@@ -0,0 +1,5 @@
+
+BACKGROUND = """
+
+
+"""