base demo

.gitignore

@@ -0,0 +1,3 @@
+venv
+.vscode
+*.pyc

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "t5_vae_flax"]
+	path = t5_vae_flax
+	url = https://github.com/Fraser-Greenlee/t5-vae-flax.git

Makefile

@@ -0,0 +1,6 @@
+
+run:
+	streamlit run streamlit_app.py
+
+test-unit:
+	streamlit hello

requirements.txt

@@ -0,0 +1,8 @@
+streamlit
+watchdog
+
+wheel
+requests
+flax
+
+-e git+https://github.com/huggingface/transformers.git#egg=transformers[flax]

streamlit_app.py

@@ -0,0 +1,95 @@
+import streamlit as st
+import jax.numpy as jnp
+from transformers import AutoTokenizer
+from transformers.models.t5.modeling_flax_t5 import shift_tokens_right
+from t5_vae_flax.src.t5_vae import FlaxT5VaeForAutoencoding
+
+
+st.title('T5-VAE')
+st.text('''
+Try interpolating between lines of Python code using this T5-VAE.
+''')
+
+
+@st.cache(allow_output_mutation=True)
+def get_model():
+    tokenizer = AutoTokenizer.from_pretrained("t5-base")
+    model = FlaxT5VaeForAutoencoding.from_pretrained("flax-community/t5-vae-python")
+    assert model.params['t5']['shared']['embedding'].shape[0] == len(tokenizer), "T5 Tokenizer doesn't match T5Vae embedding size."
+    return model, tokenizer
+
+
+model, tokenizer = get_model()
+
+
+def add_decoder_input_ids(examples):
+    arr_input_ids = jnp.array(examples["input_ids"])
+    pad = tokenizer.pad_token_id * jnp.ones((arr_input_ids.shape[0], 1), dtype=jnp.int32)
+    arr_pad_input_ids = jnp.concatenate((arr_input_ids, pad), axis=1)
+    examples['decoder_input_ids'] = shift_tokens_right(arr_pad_input_ids, tokenizer.pad_token_id, model.config.decoder_start_token_id)
+
+    arr_attention_mask = jnp.array(examples['attention_mask'])
+    ones = jnp.ones((arr_attention_mask.shape[0], 1), dtype=jnp.int32)
+    examples['decoder_attention_mask'] = jnp.concatenate((ones, arr_attention_mask), axis=1)
+
+    for k in ['decoder_input_ids', 'decoder_attention_mask']:
+        examples[k] = examples[k].tolist()
+
+    return examples
+
+
+def prepare_inputs(inputs):
+    for k, v in inputs.items():
+        inputs[k] = jnp.array(v)
+    return add_decoder_input_ids(inputs)
+
+
+def get_latent(text):
+    return model(**prepare_inputs(tokenizer([text]))).latent_codes[0]
+
+
+def tokens_from_latent(latent_codes):
+    model.config.is_encoder_decoder = True
+    output_ids = model.generate(
+        latent_codes=jnp.array([latent_codes]),
+        bos_token_id=model.config.decoder_start_token_id,
+        min_length=1,
+        max_length=32,
+    )
+    return output_ids
+
+
+def slerp(ratio, t1, t2):
+    '''
+        Perform a spherical interpolation between 2 vectors.
+        Most of the volume of a high-dimensional orange is in the skin, not the pulp.
+        This also applies for multivariate Gaussian distributions.
+        To that end we can interpolate between samples by following the surface of a n-dimensional sphere rather than a straight line.
+
+        Args:
+            ratio: Interpolation ratio.
+            t1: Tensor1
+            t2: Tensor2
+    '''
+    low_norm = t1 / jnp.linalg.norm(t1, axis=1, keepdims=True)
+    high_norm = t2 / jnp.linalg.norm(t2, axis=1, keepdims=True)
+    omega = jnp.arccos((low_norm * high_norm).sum(1))
+    so = jnp.sin(omega)
+    res = (jnp.sin((1.0 - ratio) * omega) / so)[0] * t1 + (jnp.sin(ratio * omega) / so)[0] * t2
+    return res
+
+
+def decode(ratio, txt_1, txt_2):
+    if not txt_1 or not txt_2:
+        return ''
+    lt_1, lt_2 = get_latent(txt_1), get_latent(txt_2)
+    lt_new = slerp(ratio, lt_1, lt_2)
+    tkns = tokens_from_latent(lt_new)
+    return tokenizer.decode(tkns.sequences[0], skip_special_tokens=True)
+
+
+# TODO while loop here?
+st.text_input("x = 3",          key="in_1")
+st.text_input("y += 'hello'",   key="in_2")
+r = st.slider('Interpolation Ratio')
+st.write(decode(r, st.session_state.in_1, st.session_state.in_2))

t5_vae_flax

@@ -0,0 +1 @@
+Subproject commit 0c030dca4751e6def730968a2f33fe093a608cdb