| 306 | token2id = {y: x for x, y in id2token.items()} |
| 307 | |
| 308 | class MaskedTextGenerator(keras.callbacks.Callback): |
| 309 | def __init__(self, sample_tokens, top_k=5): |
| 310 | # encoded review |
| 311 | self.sample_tokens = sample_tokens |
| 312 | self.k = top_k |
| 313 | |
| 314 | def decode(self, tokens): |
| 315 | return " ".join([id2token[t] for t in tokens if t != 0]) |
| 316 | |
| 317 | def convert_ids_to_tokens(self, id): |
| 318 | return id2token[id] |
| 319 | |
| 320 | def on_epoch_end(self, epoch, logs=None): |
| 321 | prediction = self.model.predict(self.sample_tokens) |
| 322 | # index of token2id['[mask]'] in list: |
| 323 | masked_index = np.where(self.sample_tokens == mask_token_id) |
| 324 | masked_index = masked_index[1] |
| 325 | mask_prediction = prediction[0][masked_index] |
| 326 | |
| 327 | top_indices = mask_prediction[0].argsort()[-self.k :][::-1] |
| 328 | values = mask_prediction[0][top_indices] |
| 329 | |
| 330 | for i in range(len(top_indices)): |
| 331 | p = top_indices[i] |
| 332 | v = values[i] |
| 333 | tokens = np.copy(sample_tokens[0]) |
| 334 | # fill in the blank: |
| 335 | tokens[masked_index[0]] = p |
| 336 | result = { |
| 337 | "input_text": self.decode(sample_tokens[0].numpy()), |
| 338 | "prediction": self.decode(tokens), |
| 339 | "probability": v, |
| 340 | #"predicted mask token": self.convert_ids_to_tokens(p), |
| 341 | } |
| 342 | pprint(result) |
| 343 | |
| 344 | sample_tokens = vectorize_layer(["# include < identifier . identifier > # include < identifier . [mask] > int identifier ( int identifier operator int operator identifier operator int identifier ) { }"]) |
| 345 | |