| 250 | self.decoding_params = decoding_params |
| 251 | |
| 252 | def compute(self, input_tokens, beam_width, max_seq_len, top_k=1, top_p = 0.0, beam_search_diversity_rate = 0.0, |
| 253 | temperature = 1.0, len_penalty = 0.0, repetition_penalty = 1.0, random_seed=0): |
| 254 | |
| 255 | input_ids = tf.cast(input_tokens.input_ids, tf.int32) # maybe convert to int32 |
| 256 | |
| 257 | mem_seq_len = 0 |
| 258 | if hasattr(input_tokens, "attention_mask"): |
| 259 | mem_seq_len = np.sum(input_tokens.attention_mask, axis=1) |
| 260 | else: |
| 261 | mem_seq_len = input_tokens.seq_len |
| 262 | mem_seq_len = tf.cast(mem_seq_len, tf.int32) |
| 263 | |
| 264 | encoder_outputs = ftt5_encoder(input_ids, mem_seq_len, self.encoder_params) |
| 265 | ft_decoding_output_ids, ft_decoding_seq_lens, ft_output_log_probs, ft_cum_log_probs = ftt5_decoding(encoder_outputs, |
| 266 | mem_seq_len, |
| 267 | self.decoding_params, |
| 268 | max_seq_len, |
| 269 | beam_width, |
| 270 | top_k, |
| 271 | top_p, |
| 272 | beam_search_diversity_rate, |
| 273 | temperature, |
| 274 | len_penalty, |
| 275 | repetition_penalty, |
| 276 | random_seed=random_seed) |
| 277 | |
| 278 | ft_decoding_output_ids = tf.reshape(ft_decoding_output_ids, [-1, beam_width, max_seq_len]) |
| 279 | ft_decoding_seq_lens = tf.reshape(ft_decoding_seq_lens, [-1, beam_width]) |
| 280 | |
| 281 | return ft_decoding_output_ids.numpy(), ft_decoding_seq_lens.numpy() |