Spaces:
Build error
Build error
File size: 6,329 Bytes
6a177e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 |
import logging
from typing import List
import numpy as np
import tensorflow as tf
from transformers import BertTokenizer, TFAutoModelForMaskedLM
from rhyme_with_ai.token_weighter import TokenWeighter
from rhyme_with_ai.utils import pairwise
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)
|