Spaces:
Build error
Build error
| import copy | |
| import logging | |
| from typing import List | |
| import numpy as np | |
| import tensorflow as tf | |
| import streamlit as st | |
| from transformers import BertTokenizer, TFAutoModelForMaskedLM | |
| from rhyme_with_ai.utils import color_new_words, pairwise, sanitize | |
| from rhyme_with_ai.token_weighter import TokenWeighter | |
| DEFAULT_QUERY = "Machines will take over the world soon" | |
| N_RHYMES = 10 | |
| ITER_FACTOR = 5 | |
| LANGUAGE = st.sidebar.radio("Language", ["english", "dutch"],0) | |
| if LANGUAGE == "english": | |
| MODEL_PATH = "bert-large-cased-whole-word-masking" | |
| elif LANGUAGE == "dutch": | |
| MODEL_PATH = "GroNLP/bert-base-dutch-cased" | |
| else: | |
| raise NotImplementedError(f"Unsupported language ({LANGUAGE}) expected 'english' or 'dutch'.") | |
| def main(): | |
| st.markdown( | |
| "<sup>Created with " | |
| "[Datamuse](https://www.datamuse.com/api/), " | |
| "[Mick's rijmwoordenboek](https://rijmwoordenboek.nl), " | |
| "[Hugging Face](https://huggingface.co/), " | |
| "[Streamlit](https://streamlit.io/) and " | |
| "[App Engine](https://cloud.google.com/appengine/)." | |
| " Read our [blog](https://blog.godatadriven.com/rhyme-with-ai) " | |
| "or check the " | |
| "[source](https://github.com/godatadriven/rhyme-with-ai).</sup>", | |
| unsafe_allow_html=True, | |
| ) | |
| st.title("Rhyme with AI") | |
| query = get_query() | |
| if not query: | |
| query = DEFAULT_QUERY | |
| rhyme_words_options = query_rhyme_words(query, n_rhymes=N_RHYMES,language=LANGUAGE) | |
| if rhyme_words_options: | |
| logging.getLogger(__name__).info("Got rhyme words: %s", rhyme_words_options) | |
| start_rhyming(query, rhyme_words_options) | |
| else: | |
| st.write("No rhyme words found") | |
| def get_query(): | |
| q = sanitize( | |
| st.text_input("Write your first line and press ENTER to rhyme:", DEFAULT_QUERY) | |
| ) | |
| if not q: | |
| return DEFAULT_QUERY | |
| return q | |
| def start_rhyming(query, rhyme_words_options): | |
| st.markdown("## My Suggestions:") | |
| progress_bar = st.progress(0) | |
| status_text = st.empty() | |
| max_iter = len(query.split()) * ITER_FACTOR | |
| rhyme_words = rhyme_words_options[:N_RHYMES] | |
| model, tokenizer = load_model(MODEL_PATH) | |
| sentence_generator = RhymeGenerator(model, tokenizer) | |
| sentence_generator.start(query, rhyme_words) | |
| current_sentences = [" " for _ in range(N_RHYMES)] | |
| for i in range(max_iter): | |
| previous_sentences = copy.deepcopy(current_sentences) | |
| current_sentences = sentence_generator.mutate() | |
| display_output(status_text, query, current_sentences, previous_sentences) | |
| progress_bar.progress(i / (max_iter - 1)) | |
| st.balloons() | |
| def load_model(model_path): | |
| return ( | |
| TFAutoModelForMaskedLM.from_pretrained(model_path), | |
| BertTokenizer.from_pretrained(model_path), | |
| ) | |
| def display_output(status_text, query, current_sentences, previous_sentences): | |
| print_sentences = [] | |
| for new, old in zip(current_sentences, previous_sentences): | |
| formatted = color_new_words(new, old) | |
| after_comma = "<li>" + formatted.split(",")[1][:-2] + "</li>" | |
| print_sentences.append(after_comma) | |
| status_text.markdown( | |
| query + ",<br>" + "".join(print_sentences), unsafe_allow_html=True | |
| ) | |
| class RhymeGenerator: | |
| def __init__( | |
| self, | |
| model: TFAutoModelForMaskedLM, | |
| tokenizer: BertTokenizer, | |
| token_weighter: TokenWeighter = None, | |
| ): | |
| """Generate rhymes. | |
| Parameters | |
| ---------- | |
| model : Model for masked language modelling | |
| tokenizer : Tokenizer for model | |
| token_weighter : Class that weighs tokens | |
| """ | |
| self.model = model | |
| self.tokenizer = tokenizer | |
| if token_weighter is None: | |
| token_weighter = TokenWeighter(tokenizer) | |
| self.token_weighter = token_weighter | |
| self._logger = logging.getLogger(__name__) | |
| self.tokenized_rhymes_ = None | |
| self.position_probas_ = None | |
| # Easy access. | |
| self.comma_token_id = self.tokenizer.encode(",", add_special_tokens=False)[0] | |
| self.period_token_id = self.tokenizer.encode(".", add_special_tokens=False)[0] | |
| self.mask_token_id = self.tokenizer.mask_token_id | |
| def start(self, query: str, rhyme_words: List[str]) -> None: | |
| """Start the sentence generator. | |
| Parameters | |
| ---------- | |
| query : Seed sentence | |
| rhyme_words : Rhyme words for next sentence | |
| """ | |
| # TODO: What if no content? | |
| self._logger.info("Got sentence %s", query) | |
| tokenized_rhymes = [ | |
| self._initialize_rhymes(query, rhyme_word) for rhyme_word in rhyme_words | |
| ] | |
| # Make same length. | |
| self.tokenized_rhymes_ = tf.keras.preprocessing.sequence.pad_sequences( | |
| tokenized_rhymes, padding="post", value=self.tokenizer.pad_token_id | |
| ) | |
| p = self.tokenized_rhymes_ == self.tokenizer.mask_token_id | |
| self.position_probas_ = p / p.sum(1).reshape(-1, 1) | |
| def _initialize_rhymes(self, query: str, rhyme_word: str) -> List[int]: | |
| """Initialize the rhymes. | |
| * Tokenize input | |
| * Append a comma if the sentence does not end in it (might add better predictions as it | |
| shows the two sentence parts are related) | |
| * Make second line as long as the original | |
| * Add a period | |
| Parameters | |
| ---------- | |
| query : First line | |
| rhyme_word : Last word for second line | |
| Returns | |
| ------- | |
| Tokenized rhyme lines | |
| """ | |
| query_token_ids = self.tokenizer.encode(query, add_special_tokens=False) | |
| rhyme_word_token_ids = self.tokenizer.encode( | |
| rhyme_word, add_special_tokens=False | |
| ) | |
| if query_token_ids[-1] != self.comma_token_id: | |
| query_token_ids.append(self.comma_token_id) | |
| magic_correction = len(rhyme_word_token_ids) + 1 # 1 for comma | |
| return ( | |
| query_token_ids | |
| + [self.tokenizer.mask_token_id] * (len(query_token_ids) - magic_correction) | |
| + rhyme_word_token_ids | |
| + [self.period_token_id] | |
| ) | |
| def mutate(self): | |
| """Mutate the current rhymes. | |
| Returns | |
| ------- | |
| Mutated rhymes | |
| """ | |
| self.tokenized_rhymes_ = self._mutate( | |
| self.tokenized_rhymes_, self.position_probas_, self.token_weighter.proba | |
| ) | |
| rhymes = [] | |
| for i in range(len(self.tokenized_rhymes_)): | |
| rhymes.append( | |
| self.tokenizer.convert_tokens_to_string( | |
| self.tokenizer.convert_ids_to_tokens( | |
| self.tokenized_rhymes_[i], skip_special_tokens=True | |
| ) | |
| ) | |
| ) | |
| return rhymes | |
| def _mutate( | |
| self, | |
| tokenized_rhymes: np.ndarray, | |
| position_probas: np.ndarray, | |
| token_id_probas: np.ndarray, | |
| ) -> np.ndarray: | |
| replacements = [] | |
| for i in range(tokenized_rhymes.shape[0]): | |
| mask_idx, masked_token_ids = self._mask_token( | |
| tokenized_rhymes[i], position_probas[i] | |
| ) | |
| tokenized_rhymes[i] = masked_token_ids | |
| replacements.append(mask_idx) | |
| predictions = self._predict_masked_tokens(tokenized_rhymes) | |
| for i, token_ids in enumerate(tokenized_rhymes): | |
| replace_ix = replacements[i] | |
| token_ids[replace_ix] = self._draw_replacement( | |
| predictions[i], token_id_probas, replace_ix | |
| ) | |
| tokenized_rhymes[i] = token_ids | |
| return tokenized_rhymes | |
| def _mask_token(self, token_ids, position_probas): | |
| """Mask line and return index to update.""" | |
| token_ids = self._mask_repeats(token_ids, position_probas) | |
| ix = self._locate_mask(token_ids, position_probas) | |
| token_ids[ix] = self.mask_token_id | |
| return ix, token_ids | |
| def _locate_mask(self, token_ids, position_probas): | |
| """Update masks or a random token.""" | |
| if self.mask_token_id in token_ids: | |
| # Already masks present, just return the last. | |
| # We used to return thee first but this returns worse predictions. | |
| return np.where(token_ids == self.tokenizer.mask_token_id)[0][-1] | |
| return np.random.choice(range(len(position_probas)), p=position_probas) | |
| def _mask_repeats(self, token_ids, position_probas): | |
| """Repeated tokens are generally of less quality.""" | |
| repeats = [ | |
| ii for ii, ids in enumerate(pairwise(token_ids[:-2])) if ids[0] == ids[1] | |
| ] | |
| for ii in repeats: | |
| if position_probas[ii] > 0: | |
| token_ids[ii] = self.mask_token_id | |
| if position_probas[ii + 1] > 0: | |
| token_ids[ii + 1] = self.mask_token_id | |
| return token_ids | |
| def _predict_masked_tokens(self, tokenized_rhymes): | |
| return self.model(tf.constant(tokenized_rhymes))[0] | |
| def _draw_replacement(self, predictions, token_probas, replace_ix): | |
| """Get probability, weigh and draw.""" | |
| # TODO (HG): Can't we softmax when calling the model? | |
| probas = tf.nn.softmax(predictions[replace_ix]).numpy() * token_probas | |
| probas /= probas.sum() | |
| return np.random.choice(range(len(probas)), p=probas) | |
| if __name__ == "__main__": | |
| logging.basicConfig(level=logging.INFO) | |
| main() | |