| 1415 | compute_dtype=self.compute_dtype) |
| 1416 | |
| 1417 | def encode(self, |
| 1418 | encoder_input_tokens=None, |
| 1419 | encoder_segment_ids=None, |
| 1420 | encoder_dense_inputs=None, |
| 1421 | encoder_dense_segment_ids=None, |
| 1422 | training=False): |
| 1423 | eligible_position_array = [] |
| 1424 | if encoder_input_tokens is not None: |
| 1425 | eligible_position_array.append( |
| 1426 | tf.cast(tf.not_equal(encoder_input_tokens, 0), self.compute_dtype)) |
| 1427 | if encoder_dense_inputs is not None: |
| 1428 | eligible_dense_positions = tf.cast( |
| 1429 | tf.reduce_any(tf.not_equal(encoder_dense_inputs, 0), axis=-1), |
| 1430 | self.compute_dtype) |
| 1431 | eligible_position_array.append(eligible_dense_positions) |
| 1432 | if not eligible_position_array: |
| 1433 | raise ValueError("At least one of encoder_input_tokens and" |
| 1434 | " encoder_dense_inputs must be provided.") |
| 1435 | |
| 1436 | eligible_positions = tf.concat(eligible_position_array, axis=1) |
| 1437 | encoder_mask = make_attention_mask( |
| 1438 | eligible_positions, eligible_positions, dtype=tf.bool) |
| 1439 | |
| 1440 | encoder_segment_id_array = [] |
| 1441 | if encoder_segment_ids is not None: |
| 1442 | encoder_segment_id_array.append(encoder_segment_ids) |
| 1443 | if encoder_dense_segment_ids is not None: |
| 1444 | encoder_segment_id_array.append(encoder_dense_segment_ids) |
| 1445 | if encoder_segment_id_array: |
| 1446 | encoder_segment_ids = tf.concat(encoder_segment_id_array, axis=1) |
| 1447 | segment_mask = make_attention_mask( |
| 1448 | encoder_segment_ids, encoder_segment_ids, tf.equal, dtype=tf.bool) |
| 1449 | encoder_mask = tf.math.logical_and(encoder_mask, segment_mask) |
| 1450 | encoder_mask = (1.0 - tf.cast(encoder_mask, self.compute_dtype)) * -1e9 |
| 1451 | return self.encoder( |
| 1452 | encoder_input_tokens, |
| 1453 | encoder_mask, |
| 1454 | encoder_dense_inputs, |
| 1455 | training=training) |
| 1456 | |
| 1457 | def decode( |
| 1458 | self, |