hgrif commited on
Commit
6a177e5
β€’
1 Parent(s): 5f80ca4

Add rhyme_generator.py

Browse files
app.py CHANGED
@@ -2,13 +2,12 @@ import copy
2
  import logging
3
  from typing import List
4
 
5
- import numpy as np
6
- import tensorflow as tf
7
  import streamlit as st
8
  from transformers import BertTokenizer, TFAutoModelForMaskedLM
9
- from rhyme_with_ai.utils import color_new_words, pairwise, sanitize
10
- from rhyme_with_ai.token_weighter import TokenWeighter
11
  from rhyme_with_ai.rhyme import query_rhyme_words
 
12
 
13
 
14
  DEFAULT_QUERY = "Machines will take over the world soon"
@@ -99,178 +98,6 @@ def display_output(status_text, query, current_sentences, previous_sentences):
99
  )
100
 
101
 
102
- class RhymeGenerator:
103
- def __init__(
104
- self,
105
- model: TFAutoModelForMaskedLM,
106
- tokenizer: BertTokenizer,
107
- token_weighter: TokenWeighter = None,
108
- ):
109
- """Generate rhymes.
110
-
111
- Parameters
112
- ----------
113
- model : Model for masked language modelling
114
- tokenizer : Tokenizer for model
115
- token_weighter : Class that weighs tokens
116
- """
117
-
118
- self.model = model
119
- self.tokenizer = tokenizer
120
- if token_weighter is None:
121
- token_weighter = TokenWeighter(tokenizer)
122
- self.token_weighter = token_weighter
123
- self._logger = logging.getLogger(__name__)
124
-
125
- self.tokenized_rhymes_ = None
126
- self.position_probas_ = None
127
-
128
- # Easy access.
129
- self.comma_token_id = self.tokenizer.encode(",", add_special_tokens=False)[0]
130
- self.period_token_id = self.tokenizer.encode(".", add_special_tokens=False)[0]
131
- self.mask_token_id = self.tokenizer.mask_token_id
132
-
133
- def start(self, query: str, rhyme_words: List[str]) -> None:
134
- """Start the sentence generator.
135
-
136
- Parameters
137
- ----------
138
- query : Seed sentence
139
- rhyme_words : Rhyme words for next sentence
140
- """
141
- # TODO: What if no content?
142
- self._logger.info("Got sentence %s", query)
143
- tokenized_rhymes = [
144
- self._initialize_rhymes(query, rhyme_word) for rhyme_word in rhyme_words
145
- ]
146
- # Make same length.
147
- self.tokenized_rhymes_ = tf.keras.preprocessing.sequence.pad_sequences(
148
- tokenized_rhymes, padding="post", value=self.tokenizer.pad_token_id
149
- )
150
- p = self.tokenized_rhymes_ == self.tokenizer.mask_token_id
151
- self.position_probas_ = p / p.sum(1).reshape(-1, 1)
152
-
153
- def _initialize_rhymes(self, query: str, rhyme_word: str) -> List[int]:
154
- """Initialize the rhymes.
155
-
156
- * Tokenize input
157
- * Append a comma if the sentence does not end in it (might add better predictions as it
158
- shows the two sentence parts are related)
159
- * Make second line as long as the original
160
- * Add a period
161
-
162
- Parameters
163
- ----------
164
- query : First line
165
- rhyme_word : Last word for second line
166
-
167
- Returns
168
- -------
169
- Tokenized rhyme lines
170
- """
171
-
172
- query_token_ids = self.tokenizer.encode(query, add_special_tokens=False)
173
- rhyme_word_token_ids = self.tokenizer.encode(
174
- rhyme_word, add_special_tokens=False
175
- )
176
-
177
- if query_token_ids[-1] != self.comma_token_id:
178
- query_token_ids.append(self.comma_token_id)
179
-
180
- magic_correction = len(rhyme_word_token_ids) + 1 # 1 for comma
181
- return (
182
- query_token_ids
183
- + [self.tokenizer.mask_token_id] * (len(query_token_ids) - magic_correction)
184
- + rhyme_word_token_ids
185
- + [self.period_token_id]
186
- )
187
-
188
- def mutate(self):
189
- """Mutate the current rhymes.
190
-
191
- Returns
192
- -------
193
- Mutated rhymes
194
- """
195
- self.tokenized_rhymes_ = self._mutate(
196
- self.tokenized_rhymes_, self.position_probas_, self.token_weighter.proba
197
- )
198
-
199
- rhymes = []
200
- for i in range(len(self.tokenized_rhymes_)):
201
- rhymes.append(
202
- self.tokenizer.convert_tokens_to_string(
203
- self.tokenizer.convert_ids_to_tokens(
204
- self.tokenized_rhymes_[i], skip_special_tokens=True
205
- )
206
- )
207
- )
208
- return rhymes
209
-
210
- def _mutate(
211
- self,
212
- tokenized_rhymes: np.ndarray,
213
- position_probas: np.ndarray,
214
- token_id_probas: np.ndarray,
215
- ) -> np.ndarray:
216
-
217
- replacements = []
218
- for i in range(tokenized_rhymes.shape[0]):
219
- mask_idx, masked_token_ids = self._mask_token(
220
- tokenized_rhymes[i], position_probas[i]
221
- )
222
- tokenized_rhymes[i] = masked_token_ids
223
- replacements.append(mask_idx)
224
-
225
- predictions = self._predict_masked_tokens(tokenized_rhymes)
226
-
227
- for i, token_ids in enumerate(tokenized_rhymes):
228
- replace_ix = replacements[i]
229
- token_ids[replace_ix] = self._draw_replacement(
230
- predictions[i], token_id_probas, replace_ix
231
- )
232
- tokenized_rhymes[i] = token_ids
233
-
234
- return tokenized_rhymes
235
-
236
- def _mask_token(self, token_ids, position_probas):
237
- """Mask line and return index to update."""
238
- token_ids = self._mask_repeats(token_ids, position_probas)
239
- ix = self._locate_mask(token_ids, position_probas)
240
- token_ids[ix] = self.mask_token_id
241
- return ix, token_ids
242
-
243
- def _locate_mask(self, token_ids, position_probas):
244
- """Update masks or a random token."""
245
- if self.mask_token_id in token_ids:
246
- # Already masks present, just return the last.
247
- # We used to return thee first but this returns worse predictions.
248
- return np.where(token_ids == self.tokenizer.mask_token_id)[0][-1]
249
- return np.random.choice(range(len(position_probas)), p=position_probas)
250
-
251
- def _mask_repeats(self, token_ids, position_probas):
252
- """Repeated tokens are generally of less quality."""
253
- repeats = [
254
- ii for ii, ids in enumerate(pairwise(token_ids[:-2])) if ids[0] == ids[1]
255
- ]
256
- for ii in repeats:
257
- if position_probas[ii] > 0:
258
- token_ids[ii] = self.mask_token_id
259
- if position_probas[ii + 1] > 0:
260
- token_ids[ii + 1] = self.mask_token_id
261
- return token_ids
262
-
263
- def _predict_masked_tokens(self, tokenized_rhymes):
264
- return self.model(tf.constant(tokenized_rhymes))[0]
265
-
266
- def _draw_replacement(self, predictions, token_probas, replace_ix):
267
- """Get probability, weigh and draw."""
268
- # TODO (HG): Can't we softmax when calling the model?
269
- probas = tf.nn.softmax(predictions[replace_ix]).numpy() * token_probas
270
- probas /= probas.sum()
271
- return np.random.choice(range(len(probas)), p=probas)
272
-
273
-
274
 
275
  if __name__ == "__main__":
276
  logging.basicConfig(level=logging.INFO)
 
2
  import logging
3
  from typing import List
4
 
 
 
5
  import streamlit as st
6
  from transformers import BertTokenizer, TFAutoModelForMaskedLM
7
+
8
+ from rhyme_with_ai.utils import color_new_words, sanitize
9
  from rhyme_with_ai.rhyme import query_rhyme_words
10
+ from rhyme_with_ai.rhyme_generator import RhymeGenerator
11
 
12
 
13
  DEFAULT_QUERY = "Machines will take over the world soon"
 
98
  )
99
 
100
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
101
 
102
  if __name__ == "__main__":
103
  logging.basicConfig(level=logging.INFO)
rhyme_with_ai/__init__.py ADDED
File without changes
rhyme_with_ai/rhyme_generator.py ADDED
@@ -0,0 +1,183 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import logging
2
+ from typing import List
3
+
4
+ import numpy as np
5
+ import tensorflow as tf
6
+ from transformers import BertTokenizer, TFAutoModelForMaskedLM
7
+
8
+ from rhyme_with_ai.token_weighter import TokenWeighter
9
+ from rhyme_with_ai.utils import pairwise
10
+
11
+
12
+ class RhymeGenerator:
13
+ def __init__(
14
+ self,
15
+ model: TFAutoModelForMaskedLM,
16
+ tokenizer: BertTokenizer,
17
+ token_weighter: TokenWeighter = None,
18
+ ):
19
+ """Generate rhymes.
20
+
21
+ Parameters
22
+ ----------
23
+ model : Model for masked language modelling
24
+ tokenizer : Tokenizer for model
25
+ token_weighter : Class that weighs tokens
26
+ """
27
+
28
+ self.model = model
29
+ self.tokenizer = tokenizer
30
+ if token_weighter is None:
31
+ token_weighter = TokenWeighter(tokenizer)
32
+ self.token_weighter = token_weighter
33
+ self._logger = logging.getLogger(__name__)
34
+
35
+ self.tokenized_rhymes_ = None
36
+ self.position_probas_ = None
37
+
38
+ # Easy access.
39
+ self.comma_token_id = self.tokenizer.encode(",", add_special_tokens=False)[0]
40
+ self.period_token_id = self.tokenizer.encode(".", add_special_tokens=False)[0]
41
+ self.mask_token_id = self.tokenizer.mask_token_id
42
+
43
+ def start(self, query: str, rhyme_words: List[str]) -> None:
44
+ """Start the sentence generator.
45
+
46
+ Parameters
47
+ ----------
48
+ query : Seed sentence
49
+ rhyme_words : Rhyme words for next sentence
50
+ """
51
+ # TODO: What if no content?
52
+ self._logger.info("Got sentence %s", query)
53
+ tokenized_rhymes = [
54
+ self._initialize_rhymes(query, rhyme_word) for rhyme_word in rhyme_words
55
+ ]
56
+ # Make same length.
57
+ self.tokenized_rhymes_ = tf.keras.preprocessing.sequence.pad_sequences(
58
+ tokenized_rhymes, padding="post", value=self.tokenizer.pad_token_id
59
+ )
60
+ p = self.tokenized_rhymes_ == self.tokenizer.mask_token_id
61
+ self.position_probas_ = p / p.sum(1).reshape(-1, 1)
62
+
63
+ def _initialize_rhymes(self, query: str, rhyme_word: str) -> List[int]:
64
+ """Initialize the rhymes.
65
+
66
+ * Tokenize input
67
+ * Append a comma if the sentence does not end in it (might add better predictions as it
68
+ shows the two sentence parts are related)
69
+ * Make second line as long as the original
70
+ * Add a period
71
+
72
+ Parameters
73
+ ----------
74
+ query : First line
75
+ rhyme_word : Last word for second line
76
+
77
+ Returns
78
+ -------
79
+ Tokenized rhyme lines
80
+ """
81
+
82
+ query_token_ids = self.tokenizer.encode(query, add_special_tokens=False)
83
+ rhyme_word_token_ids = self.tokenizer.encode(
84
+ rhyme_word, add_special_tokens=False
85
+ )
86
+
87
+ if query_token_ids[-1] != self.comma_token_id:
88
+ query_token_ids.append(self.comma_token_id)
89
+
90
+ magic_correction = len(rhyme_word_token_ids) + 1 # 1 for comma
91
+ return (
92
+ query_token_ids
93
+ + [self.tokenizer.mask_token_id] * (len(query_token_ids) - magic_correction)
94
+ + rhyme_word_token_ids
95
+ + [self.period_token_id]
96
+ )
97
+
98
+ def mutate(self):
99
+ """Mutate the current rhymes.
100
+
101
+ Returns
102
+ -------
103
+ Mutated rhymes
104
+ """
105
+ self.tokenized_rhymes_ = self._mutate(
106
+ self.tokenized_rhymes_, self.position_probas_, self.token_weighter.proba
107
+ )
108
+
109
+ rhymes = []
110
+ for i in range(len(self.tokenized_rhymes_)):
111
+ rhymes.append(
112
+ self.tokenizer.convert_tokens_to_string(
113
+ self.tokenizer.convert_ids_to_tokens(
114
+ self.tokenized_rhymes_[i], skip_special_tokens=True
115
+ )
116
+ )
117
+ )
118
+ return rhymes
119
+
120
+ def _mutate(
121
+ self,
122
+ tokenized_rhymes: np.ndarray,
123
+ position_probas: np.ndarray,
124
+ token_id_probas: np.ndarray,
125
+ ) -> np.ndarray:
126
+
127
+ replacements = []
128
+ for i in range(tokenized_rhymes.shape[0]):
129
+ mask_idx, masked_token_ids = self._mask_token(
130
+ tokenized_rhymes[i], position_probas[i]
131
+ )
132
+ tokenized_rhymes[i] = masked_token_ids
133
+ replacements.append(mask_idx)
134
+
135
+ predictions = self._predict_masked_tokens(tokenized_rhymes)
136
+
137
+ for i, token_ids in enumerate(tokenized_rhymes):
138
+ replace_ix = replacements[i]
139
+ token_ids[replace_ix] = self._draw_replacement(
140
+ predictions[i], token_id_probas, replace_ix
141
+ )
142
+ tokenized_rhymes[i] = token_ids
143
+
144
+ return tokenized_rhymes
145
+
146
+ def _mask_token(self, token_ids, position_probas):
147
+ """Mask line and return index to update."""
148
+ token_ids = self._mask_repeats(token_ids, position_probas)
149
+ ix = self._locate_mask(token_ids, position_probas)
150
+ token_ids[ix] = self.mask_token_id
151
+ return ix, token_ids
152
+
153
+ def _locate_mask(self, token_ids, position_probas):
154
+ """Update masks or a random token."""
155
+ if self.mask_token_id in token_ids:
156
+ # Already masks present, just return the last.
157
+ # We used to return thee first but this returns worse predictions.
158
+ return np.where(token_ids == self.tokenizer.mask_token_id)[0][-1]
159
+ return np.random.choice(range(len(position_probas)), p=position_probas)
160
+
161
+ def _mask_repeats(self, token_ids, position_probas):
162
+ """Repeated tokens are generally of less quality."""
163
+ repeats = [
164
+ ii for ii, ids in enumerate(pairwise(token_ids[:-2])) if ids[0] == ids[1]
165
+ ]
166
+ for ii in repeats:
167
+ if position_probas[ii] > 0:
168
+ token_ids[ii] = self.mask_token_id
169
+ if position_probas[ii + 1] > 0:
170
+ token_ids[ii + 1] = self.mask_token_id
171
+ return token_ids
172
+
173
+ def _predict_masked_tokens(self, tokenized_rhymes):
174
+ return self.model(tf.constant(tokenized_rhymes))[0]
175
+
176
+ def _draw_replacement(self, predictions, token_probas, replace_ix):
177
+ """Get probability, weigh and draw."""
178
+ # TODO (HG): Can't we softmax when calling the model?
179
+ probas = tf.nn.softmax(predictions[replace_ix]).numpy() * token_probas
180
+ probas /= probas.sum()
181
+ return np.random.choice(range(len(probas)), p=probas)
182
+
183
+