| import tensorflow as tf | |
| import numpy as np | |
| from models.transformer import Transformer | |
| from models.encoder import Encoder | |
| from models.decoder import Decoder | |
| from models.layers import EncoderLayer, DecoderLayer, MultiHeadAttention, point_wise_feed_forward_network | |
| from models.utils import masked_loss, masked_accuracy | |
| class Translator(tf.Module): | |
| def __init__(self, en_tokenizer, vi_tokenizer, transformer, max_tokens=512): | |
| self.tokenizers_en = en_tokenizer | |
| self.tokenizers_vi = vi_tokenizer | |
| self.transformer = transformer | |
| self.max_tokens = max_tokens | |
| def __call__(self, sentence, max_length=None): | |
| if max_length is None: | |
| max_length = self.max_tokens | |
| sentence = self.tokenizers_en.texts_to_sequences([sentence]) | |
| sentence = sentence[0] + np.zeros(self.max_tokens - len(sentence[0]), | |
| dtype=np.int32).tolist() | |
| sentence = tf.convert_to_tensor([sentence]) | |
| encoder_input = sentence | |
| start = self.tokenizers_vi.texts_to_sequences(["<sos>"])[0] | |
| end = self.tokenizers_vi.texts_to_sequences(["<eos>"])[0] | |
| output_array = tf.TensorArray(dtype=tf.int32, size=0, dynamic_size=True) | |
| output_array = output_array.write(0, start) | |
| for i in tf.range(max_length): | |
| output = tf.transpose(output_array.stack()) | |
| predictions = self.transformer([encoder_input, output], training=False) | |
| predictions = predictions[:, -1:, :] | |
| predicted_id = tf.argmax(predictions, axis=-1, output_type=tf.int32) | |
| output_array = output_array.write(i+1, predicted_id[0]) | |
| if predicted_id == end: | |
| break | |
| output = tf.transpose(output_array.stack()) | |
| text = self.tokenizers_vi.sequences_to_texts(output.numpy().tolist())[0] | |
| return text |