def test_data_generator_with_increasing_batch_size(model_data: RasaModelData): epochs = 2 data_generator = RasaBatchDataGenerator( model_data, batch_size=[1, 2], epochs=epochs, batch_strategy="balanced", shuffle=True, ) expected_batch_sizes = [[1, 1, 1, 1, 1], [2, 2, 1]] for _epoch in range(epochs): iterator = iter(data_generator) assert len(data_generator) == len(expected_batch_sizes[_epoch]) for i in range(len(data_generator)): batch, _ = next(iterator) assert len(batch) == 11 assert len(batch[0]) == expected_batch_sizes[_epoch][i] with pytest.raises(StopIteration): next(iterator) data_generator.on_epoch_end()
def rasa_predict(self, model_data: RasaModelData) -> Dict[Text, tf.Tensor]: """Custom prediction method that builds tf graph on the first call. Args: model_data: The model data to use for prediction. Return: Prediction output. """ self._training = False if not self.prepared_for_prediction: # in case the model is used for prediction without loading, e.g. directly # after training, we need to prepare the model for prediction once self.prepare_for_predict() self.prepared_for_prediction = True batch_in = RasaBatchDataGenerator.prepare_batch(model_data.data) if self._run_eagerly: return tf_utils.to_numpy_or_python_type(self.predict_step(batch_in)) if self._tf_predict_step is None: self._tf_predict_step = tf.function( self.predict_step, input_signature=self._dynamic_signature(batch_in) ) return tf_utils.to_numpy_or_python_type(self._tf_predict_step(batch_in))
def create_data_generators( model_data: RasaModelData, batch_sizes: Union[int, List[int]], epochs: int, batch_strategy: Text = SEQUENCE, eval_num_examples: int = 0, random_seed: Optional[int] = None, shuffle: bool = True, ) -> Tuple[RasaBatchDataGenerator, Optional[RasaBatchDataGenerator]]: """Create data generators for train and optional validation data. Args: model_data: The model data to use. batch_sizes: The batch size(s). epochs: The number of epochs to train. batch_strategy: The batch strategy to use. eval_num_examples: Number of examples to use for validation data. random_seed: The random seed. shuffle: Whether to shuffle data inside the data generator. Returns: The training data generator and optional validation data generator. """ validation_data_generator = None if eval_num_examples > 0: model_data, evaluation_model_data = model_data.split( eval_num_examples, random_seed, ) validation_data_generator = RasaBatchDataGenerator( evaluation_model_data, batch_size=batch_sizes, epochs=epochs, batch_strategy=batch_strategy, shuffle=shuffle, ) data_generator = RasaBatchDataGenerator( model_data, batch_size=batch_sizes, epochs=epochs, batch_strategy=batch_strategy, shuffle=shuffle, ) return data_generator, validation_data_generator
def _compile_and_fit( self, data_example: Dict[Text, Dict[Text, List[FeatureArray]]] ) -> None: """Compiles modified model and fits a sample data on it. Args: data_example: a data example that is stored with the ML component. """ self.compile(optimizer=tf.keras.optimizers.Adam(self.config[LEARNING_RATE])) label_key = LABEL_KEY if self.config[INTENT_CLASSIFICATION] else None label_sub_key = LABEL_SUB_KEY if self.config[INTENT_CLASSIFICATION] else None model_data = RasaModelData( label_key=label_key, label_sub_key=label_sub_key, data=data_example ) self._update_data_signatures(model_data) data_generator = RasaBatchDataGenerator(model_data, batch_size=1) self.fit(data_generator, verbose=False)
def load( cls, model_file_name: Text, model_data_example: RasaModelData, predict_data_example: Optional[RasaModelData] = None, finetune_mode: bool = False, *args: Any, **kwargs: Any, ) -> "RasaModel": """Loads a model from the given weights. Args: model_file_name: Path to file containing model weights. model_data_example: Example data point to construct the model architecture. predict_data_example: Example data point to speed up prediction during inference. finetune_mode: Indicates whether to load the model for further finetuning. *args: Any other non key-worded arguments. **kwargs: Any other key-worded arguments. Returns: Loaded model with weights appropriately set. """ logger.debug( f"Loading the model from {model_file_name} " f"with finetune_mode={finetune_mode}..." ) # create empty model model = cls(*args, **kwargs) learning_rate = kwargs.get("config", {}).get(LEARNING_RATE, 0.001) # need to train on 1 example to build weights of the correct size model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate)) data_generator = RasaBatchDataGenerator(model_data_example, batch_size=1) model.fit(data_generator, verbose=False) # load trained weights model.load_weights(model_file_name) # predict on one data example to speed up prediction during inference # the first prediction always takes a bit longer to trace tf function if not finetune_mode and predict_data_example: model.run_inference(predict_data_example) logger.debug("Finished loading the model.") return model
def test_data_generator_with_fixed_batch_size(model_data: RasaModelData): data_generator = RasaBatchDataGenerator(model_data, batch_size=2, epochs=1, batch_strategy="balanced", shuffle=True) expected_batch_sizes = [2, 2, 1] iterator = iter(data_generator) assert len(data_generator) == len(expected_batch_sizes) for i in range(len(data_generator)): batch, _ = next(iterator) assert len(batch) == 11 assert len(batch[0]) == expected_batch_sizes[i] with pytest.raises(StopIteration): next(iterator)
async def test_gen_batch(self, trained_policy: TEDPolicy, default_domain: Domain, stories_path: Path): training_trackers = await tests.core.test_policies.train_trackers( default_domain, stories_path, augmentation_factor=0) interpreter = RegexInterpreter() training_data, label_ids, entity_tags = trained_policy._featurize_for_training( training_trackers, default_domain, interpreter) label_data, all_labels = trained_policy._create_label_data( default_domain, interpreter) model_data = trained_policy._create_model_data(training_data, label_ids, entity_tags, all_labels) batch_size = 2 data_generator = RasaBatchDataGenerator(model_data, batch_size=batch_size, shuffle=False, batch_strategy="sequence") iterator = iter(data_generator) # model data keys were sorted, so the order is alphabetical ( ( batch_action_name_mask, batch_action_name_sentence_indices, batch_action_name_sentence_data, batch_action_name_sentence_shape, batch_dialogue_length, batch_entities_mask, batch_entities_sentence_indices, batch_entities_sentence_data, batch_entities_sentence_shape, batch_intent_mask, batch_intent_sentence_indices, batch_intent_sentence_data, batch_intent_sentence_shape, batch_label_ids, batch_slots_mask, batch_slots_sentence_indices, batch_slots_sentence_data, batch_slots_sentence_shape, ), _, ) = next(iterator) assert (batch_label_ids.shape[0] == batch_size and batch_dialogue_length.shape[0] == batch_size) # batch and dialogue dimensions are NOT combined for masks assert (batch_slots_mask.shape[0] == batch_size and batch_intent_mask.shape[0] == batch_size and batch_entities_mask.shape[0] == batch_size and batch_action_name_mask.shape[0] == batch_size) # some features might be "fake" so there sequence is `0` seq_len = max([ batch_intent_sentence_shape[1], batch_action_name_sentence_shape[1], batch_entities_sentence_shape[1], batch_slots_sentence_shape[1], ]) assert (batch_intent_sentence_shape[1] == seq_len or batch_intent_sentence_shape[1] == 0) assert (batch_action_name_sentence_shape[1] == seq_len or batch_action_name_sentence_shape[1] == 0) assert (batch_entities_sentence_shape[1] == seq_len or batch_entities_sentence_shape[1] == 0) assert (batch_slots_sentence_shape[1] == seq_len or batch_slots_sentence_shape[1] == 0) data_generator = RasaBatchDataGenerator(model_data, batch_size=batch_size, shuffle=True, batch_strategy="balanced") iterator = iter(data_generator) ( ( batch_action_name_mask, batch_action_name_sentence_indices, batch_action_name_sentence_data, batch_action_name_sentence_shape, batch_dialogue_length, batch_entities_mask, batch_entities_sentence_indices, batch_entities_sentence_data, batch_entities_sentence_shape, batch_intent_mask, batch_intent_sentence_indices, batch_intent_sentence_data, batch_intent_sentence_shape, batch_label_ids, batch_slots_mask, batch_slots_sentence_indices, batch_slots_sentence_data, batch_slots_sentence_shape, ), _, ) = next(iterator) assert (batch_label_ids.shape[0] == batch_size and batch_dialogue_length.shape[0] == batch_size) # some features might be "fake" so there sequence is `0` seq_len = max([ batch_intent_sentence_shape[1], batch_action_name_sentence_shape[1], batch_entities_sentence_shape[1], batch_slots_sentence_shape[1], ]) assert (batch_intent_sentence_shape[1] == seq_len or batch_intent_sentence_shape[1] == 0) assert (batch_action_name_sentence_shape[1] == seq_len or batch_action_name_sentence_shape[1] == 0) assert (batch_entities_sentence_shape[1] == seq_len or batch_entities_sentence_shape[1] == 0) assert (batch_slots_sentence_shape[1] == seq_len or batch_slots_sentence_shape[1] == 0)