Spaces:
Runtime error
Runtime error
| 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_alt.src.t5_vae import FlaxT5VaeForAutoencoding | |
| st.set_page_config( | |
| page_title="T5-VAE", | |
| page_icon="😐", | |
| layout="wide", | |
| initial_sidebar_state="expanded" | |
| ) | |
| st.title('T5-VAE 🙁😐🙂') | |
| st.markdown(''' | |
| 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.markdown(''' | |
| ### [t5-vae-python](https://huggingface.co/flax-community/t5-vae-python) | |
| This model is trained on lines of Python code from GitHub ([dataset](https://huggingface.co/datasets/Fraser/python-lines)). | |
| ''') | |
| 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(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) | |
| in_1 = st.text_input("A line of Python code.", "x = a - 1") | |
| in_2 = st.text_input("Another line of Python code.", "x = a + 10 * 2") | |
| r = st.slider('Python 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('Output: ' + out) | |
| st.markdown(''' | |
| ### [t5-vae-wiki](https://huggingface.co/flax-community/t5-vae-wiki) | |
| This model is trained on just 5% of the sentences on wikipedia. | |
| We'll release another model trained on the full [dataset](https://github.com/ChunyuanLI/Optimus/blob/master/download_datasets.md) soon. | |
| ''') | |
| 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('English 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('Output: ' + out) | |
| st.markdown(''' | |
| Try arithmetic in latent space. | |
| Here latent codes for each sentence are found and arithmetic is done with them. | |
| Here it runs the sum `C + (B - A) = ?` | |
| ''') | |
| 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", "A girl makes a silly face.") | |
| in_b = st.text_input("B", "Two girls are playing soccer.") | |
| in_c = st.text_input("C", "A girl is looking through a microscope.") | |
| st.markdown(''' | |
| 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('Output: ' + out) | |