rasa-pro 3.11.0__py3-none-any.whl → 3.11.0a1__py3-none-any.whl

rasa/llm_fine_tuning/notebooks/unsloth_finetuning.ipynb

@@ -0,0 +1,407 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Fine-tuning a Hugging Face base model using Unsloth and TRL\n",
+    "\n",
+    "This is a worked example of how to efficiently fine-tune a base language model from [Hugging Face Hub](https://huggingface.co/models) using the [Unsloth](https://docs.unsloth.ai) and [TRL](https://huggingface.co/docs/trl/en/index) libraries on an instruction-based task.\n",
+    "\n",
+    "Unsloth integrates with TRL in order to reduce the time and GPU memory required to fine-tune LLMs, when compared to using TRL exclusively.\n",
+    "\n",
+    "To run fine-tuning, you must have first [generated the dataset](https://rasa.com/rasa-pro/docs/operating/fine-tuning-recipe) files `train.jsonl` and `val.jsonl`, which must be in the [TRL instruction format](https://huggingface.co/docs/trl/en/sft_trainer#dataset-format-support)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Configure fine-tuning environment\n",
+    "\n",
+    "This notebook has been tested on on a [GCP Vertex AI Workbench instance](https://cloud.google.com/vertex-ai/docs/workbench/instances/introduction) with machine type `a2-highgpu-1g` in the `asia-southeast1-b` zone, which had the following hardware:\n",
+    "- Single NVIDIA A100 GPU (40GB VRAM)\n",
+    "- 12 core CPU with 85B RAM\n",
+    "- 256GB disk\n",
+    "\n",
+    "It has also been tested on on a [AWS Sagemaker Notebook instance](https://docs.aws.amazon.com/sagemaker/latest/dg/nbi.html) with machine type `ml.g4dn.4xlarge` in the `eu-central-1` region, which had the following hardware:\n",
+    "- Single NVIDIA T4 GPU (16GB VRAM)\n",
+    "- 16 core CPU with 64GB RAM\n",
+    "- 256GB disk\n",
+    "\n",
+    "In both setups, the notebook was executed in a Linux environment with the following software already installed:\n",
+    "- Python 3.10\n",
+    "- CUDA Toolkit 12.1\n",
+    "- PyTorch 2.2\n",
+    "\n",
+    "In order to run fine-tuning yourself, upload this notebook along with your dataset files to your own environment.\n",
+    "\n",
+    "It is highly recommended that you restart your notebook kernel and re-run all notebook cells every time you wish to perform fine-tuning.\n",
+    "\n",
+    "> Despite the fact that this notebook will work with a relatively underpowered GPU, such as the NVIDIA T4, running LLM fine-tuning and inference will be very slow.\n",
+    ">\n",
+    "> It is highly recommended that you use an NVIDIA A100 or other similar GPU types.\n",
+    ">\n",
+    "> The code presented here has been configured for use with an NVIDIA A100, please take note of later comments on the changes you should make to the code when using a different type of GPU.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Install Python requirements\n",
+    "\n",
+    "The following `pip` commands will install Unsloth and other required Python packages when [Conda](https://anaconda.org/anaconda/conda) is used to manage your Python environment, which is the case for many Jupyter notebook runtimes.\n",
+    "\n",
+    "If Conda is not used in your environment, please follow [these alternative instructions](https://github.com/unslothai/unsloth?tab=readme-ov-file#pip-installation) for installing Unsloth."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%sh\n",
+    "# install unsloth and other dependencies\n",
+    "pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\"\n",
+    "pip install --no-deps \"xformers<=0.0.26\" \"trl<0.9.0\" peft accelerate bitsandbytes huggingface_hub[cli]\n",
+    "# remove tpu-only package that is installed by default on gcp runtimes, even when only using gpu\n",
+    "pip uninstall torch-xla -y"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Download base model\n",
+    "\n",
+    "You can download the model you want to fine-tune from Hugging Face Hub using the [official CLI](https://huggingface.co/docs/huggingface_hub/en/guides/cli) with an [API access token](https://huggingface.co/docs/transformers.js/en/guides/private#step-1-generating-a-user-access-token) as per the code below. Make sure you first update the `HUGGINGFACE_TOKEN` and `BASE_MODEL` environment variables with your own values.\n",
+    "\n",
+    "When testing this notebook, the [Llama 3.1 8B Instruct](https://huggingface.co/meta-llama/Meta-Llama-3.1-8B-Instruct) model was used. Note that `meta-llama/Meta-Llama-3.1-8B-Instruct` is a [gated model](https://huggingface.co/docs/hub/en/models-gated) that you must first request access to before using. \n",
+    "\n",
+    "You can use any other PyTorch model available on [Hugging Face Hub](https://huggingface.co/models). It is recommended that you use a model that has been pre-trained on instructional tasks, such as the [CodeLlama 13B Instruct](https://huggingface.co/codellama/CodeLlama-13b-Instruct-hf) model.\n",
+    "\n",
+    "Pre-trained models with more parameters will generally perform better at tasks than models with fewer parameters. However, the size of model you can use is limited by how much memory your GPU has.\n",
+    "\n",
+    "Alternatively, if you already have a PyTorch model directory to hand, you can upload it to your notebook environment manually."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%sh\n",
+    "# TODO: update with your values\n",
+    "export HUGGINGFACE_TOKEN=\"CHANGEME\"\n",
+    "export BASE_MODEL=\"meta-llama/Meta-Llama-3.1-8B-Instruct\"\n",
+    "\n",
+    "# download model\n",
+    "huggingface-cli download \"${BASE_MODEL}\" \\\n",
+    "    --token \"${HUGGINGFACE_TOKEN}\" \\\n",
+    "    --local-dir \"./base_model\" \\\n",
+    "    --exclude \"*.bin*\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Load and quantize base model\n",
+    "\n",
+    "The [quantization of model parameters](https://huggingface.co/docs/optimum/en/concept_guides/quantization) can significantly reduce the GPU memory required to run model fine-tuning and inference, at the cost of model accuracy.\n",
+    "\n",
+    "Here, the base model is loaded from disk and quantized into an 8-bit representation on the fly using the [BitsAndBytes](https://huggingface.co/docs/transformers/main/en/quantization/bitsandbytes) library.\n",
+    "\n",
+    "If you are using a GPU with relatively little memory, such as the NVIDIA T4, or if you are using a base model larger than `meta-llama/Meta-Llama-3.1-8B-Instruct`, you may be required to use 4-bit quantization (e.g. `load_in_4bit = True`) in order to avoid \n",
+    "out of memory (OOM) errors."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import BitsAndBytesConfig\n",
+    "from unsloth import FastLanguageModel\n",
+    "\n",
+    "# configure quantization method for base model\n",
+    "quantization_config = BitsAndBytesConfig(\n",
+    "    load_in_8bit=True,\n",
+    ")\n",
+    "\n",
+    "# load quantized model and tokenizer from disk\n",
+    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    model_name=\"./base_model\",\n",
+    "    max_seq_length=2048,\n",
+    "    quantization_config=quantization_config,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Load training and validation datasets\n",
+    "\n",
+    "The following code loads the training and validation datasets from the `train.jsonl` and `val.json` files, respectively\n",
+    "\n",
+    "As the files use the TRL instruction format, the TRL trainer used later will be able to [automatically parse](https://huggingface.co/docs/trl/en/sft_trainer#dataset-format-support) the datasets and [generate the prompts from a template](https://huggingface.co/docs/transformers/en/chat_templating) configured in the tokenizer.\n",
+    "\n",
+    "Prompt templates vary between models and TRL will infer the correct template from your base model. If this is not available in your base model or if you wish to change it, you can set your own [template string](https://huggingface.co/docs/transformers/en/chat_templating#advanced-adding-and-editing-chat-templates).\n",
+    "\n",
+    "You can also define your own [prompt formatting function](https://huggingface.co/docs/trl/en/sft_trainer#format-your-input-prompts) in order to have full control of how the prompts are constructed."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import datasets\n",
+    "from trl.extras.dataset_formatting import get_formatting_func_from_dataset\n",
+    "\n",
+    "# load datasets from disk\n",
+    "train_dataset = datasets.load_dataset(\n",
+    "    \"json\", data_files={\"train\": \"train.jsonl\"}, split=\"train\"\n",
+    ")\n",
+    "eval_dataset = datasets.load_dataset(\n",
+    "    \"json\", data_files={\"eval\": \"val.jsonl\"}, split=\"eval\"\n",
+    ")\n",
+    "\n",
+    "# test prompt templating on example from dataset\n",
+    "print(get_formatting_func_from_dataset(train_dataset, tokenizer)(eval_dataset[0]))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Configure trainer\n",
+    "\n",
+    "Below, the arguments for the supervised fine-tuning (SFT) trainer are configured. Their values were chosen somewhat arbitrarily and resulted in satisfactory results during testing.\n",
+    "\n",
+    "It is recommended that you read the official documentation and experiment with the arguments passed to `SFTConfig` (see [here](https://huggingface.co/docs/trl/main/en/sft_trainer#trl.SFTTrainer)) and `SFTTrainer` (see [here](https://huggingface.co/docs/trl/main/en/sft_trainer#trl.SFTTrainer)).\n",
+    "\n",
+    "For example:\n",
+    "- If you get an OOM error when running fine-tuning, you can reduce `per_device_train_batch_size` in order to reduce the memory footprint. However, if your GPU has sufficient memory, you can try increasing it in order to reduce the total number of training steps.\n",
+    "- Consider setting `max_steps`, as you may not need to perform all epochs in order to achieve optimal model accuracy. Conversely, you may see better model accuracy by increasing `num_train_epochs`.\n",
+    "- If fine-tuning is taking too long, you can increase `eval_steps` in order to reduce how often validation is performed. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from transformers import TrainingArguments\n",
+    "from trl import SFTTrainer\n",
+    "\n",
+    "# configure training args\n",
+    "args = TrainingArguments(\n",
+    "    # training\n",
+    "    per_device_train_batch_size=8,\n",
+    "    warmup_steps=50,\n",
+    "    num_train_epochs=4,\n",
+    "    learning_rate=0.0001,\n",
+    "    lr_scheduler_type=\"cosine\",\n",
+    "    optim=\"adamw_bnb_8bit\",\n",
+    "    weight_decay=0.0,\n",
+    "    logging_steps=1,\n",
+    "    # datatypes\n",
+    "    fp16=not torch.cuda.is_bf16_supported(),\n",
+    "    bf16=torch.cuda.is_bf16_supported(),\n",
+    "    # evaluation\n",
+    "    eval_strategy=\"steps\",\n",
+    "    eval_steps=50,\n",
+    "    per_device_eval_batch_size=8,\n",
+    "    output_dir=\"outputs\",\n",
+    ")\n",
+    "\n",
+    "# setup trainer\n",
+    "trainer = SFTTrainer(\n",
+    "    model=model,\n",
+    "    tokenizer=tokenizer,\n",
+    "    train_dataset=train_dataset,\n",
+    "    eval_dataset=eval_dataset,\n",
+    "    args=args,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. Perform supervised fine-tuning\n",
+    "\n",
+    "In the code below, fine-tuning is performed using the previously congfigured trainer.\n",
+    "\n",
+    "When testing this step on an NVIDIA A100 using the configuration defined above, it took around 12 minutes to perform fine-tuning with a training dataset containing around 500 examples.\n",
+    "\n",
+    "After fine-tuning, the base model and fine-tuned adapters are [merged together and saved to disk](https://docs.unsloth.ai/basics/saving-models/saving-to-vllm) in 16-bit for future compatibility with the [vLLM](https://github.com/vllm-project/vllm) model serving library.\n",
+    "\n",
+    "If you are using a relatively small GPU, such as the NVIDIA T4, you may have to save the model in 4-bit instead (e.g. `save_method = \"merged_4bit_forced\"`)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# run fine-tuning\n",
+    "finetune_metrics = trainer.train()\n",
+    "\n",
+    "# save model to disk in 16-bit\n",
+    "model.save_pretrained_merged(\n",
+    "    \"./finetune_model\", tokenizer, save_method=\"merged_16bit\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 8. Visualize fine-tuning metrics\n",
+    "\n",
+    "Some of the metrics collected during fine-tuning are visualised below in order for you to diagnose any potential issues with the model.\n",
+    "\n",
+    "Specifically, the training and validation losses are plotted against the training step number. Please check the plot for the following:\n",
+    "- Ideally, as the fine-tuning steps increase, the training and validation losses should decrease and converge. \n",
+    "- If both loss curves do not converge, it may be worth performing more fine-tuning steps or epochs. This is known as [underfitting](https://www.ibm.com/topics/underfitting).\n",
+    "- If the validation loss suddenly starts to increase while the training loss continues to decrease or converge, you should decrease your total number of steps or epochs. This is known as [overfitting](https://www.ibm.com/topics/overfitting)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# plot step against train and val losses\n",
+    "fig, ax = plt.subplots()\n",
+    "log_history = pd.DataFrame(trainer.state.log_history)\n",
+    "eval_loss = (\n",
+    "    log_history[[\"step\", \"eval_loss\"]]\n",
+    "    .dropna()\n",
+    "    .plot(x=\"step\", ax=ax)\n",
+    ")\n",
+    "train_loss = (\n",
+    "    log_history[[\"step\", \"loss\"]]\n",
+    "    .dropna()\n",
+    "    .plot(x=\"step\", ax=ax)\n",
+    ")\n",
+    "fig.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 9. Run ad hoc inference\n",
+    "\n",
+    "You can load your fine-tuned model from disk using Unsloth and use it to run optimized inference on individual inputs of your choosing using the code below.\n",
+    "\n",
+    "Note that the inputs passed to model are in the [TRL convertsational format](https://huggingface.co/docs/trl/en/sft_trainer#dataset-format-support) as the Hugging Face [chat template requires them to be](https://huggingface.co/docs/transformers/main/en/chat_templating#how-do-i-use-chat-templates). During training TRL will [automatically convert the instruction format to the conversational format](https://github.com/huggingface/trl/blob/main/trl/extras/dataset_formatting.py). However, you have to do this yourself when applying chat templates manually for inference."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import TextStreamer\n",
+    "from unsloth import FastLanguageModel\n",
+    "\n",
+    "model, tokenizer = FastLanguageModel.from_pretrained(\"./finetune_model\")\n",
+    "FastLanguageModel.for_inference(model)  # enable inference optimizations\n",
+    "streamer = TextStreamer(tokenizer)  # stream model outputs as they are generated\n",
+    "\n",
+    "# the content to include in the input prompt\n",
+    "# by default, a value from the validation dataset as example\n",
+    "content = eval_dataset[\"prompt\"][0]\n",
+    "\n",
+    "# apply prompt template and tokenize\n",
+    "input_ids = tokenizer.apply_chat_template(\n",
+    "    [{\"role\": \"user\", \"content\": content}],  # in the TRL conversational format\n",
+    "    tokenize=True,\n",
+    "    add_generation_prompt=True,\n",
+    "    return_tensors=\"pt\",\n",
+    ").to(\"cuda\")\n",
+    "\n",
+    "# generate model output from user input\n",
+    "_ = model.generate(\n",
+    "    input_ids=input_ids,\n",
+    "    streamer=streamer,  # remove streamer if you want whole output at end\n",
+    "    max_new_tokens=64,  # set the limit on how many tokens are generated\n",
+    "    do_sample=False,    # disable random sampling for deterministic outputs\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 10. Export fine-tuned model\n",
+    "\n",
+    "Lastly, export your fine-tuned model directory to an appropriate storage location that can be easily accessed later for [deployment](https://rasa.com/rasa-pro/docs/building-assistants/self-hosted-llm).\n",
+    "\n",
+    "It is recommended that you use a cloud object store, such as [Amazon S3](https://aws.amazon.com/s3/) or [Google Cloud Storage](https://cloud.google.com/storage).\n",
+    "\n",
+    "Uncomment and run the corresponding commands below for your cloud provider, making sure to first update the environment variables with your own values. It is assumed that:\n",
+    "- your bucket already exists\n",
+    "- you have already installed the CLI tool for your cloud provider\n",
+    "- you have already authenticated with your cloud provider and have sufficient permissions to write to your bucket"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%sh\n",
+    "export LOCAL_MODEL_PATH=\"./finetune_model\"\n",
+    "\n",
+    "# if using amazon\n",
+    "# export S3_MODEL_URI=\"s3://CHANGEME\" # update with your value\n",
+    "# aws s3 cp \"${LOCAL_MODEL_PATH}\" \"${S3_MODEL_URI}\" --recursive\n",
+    "\n",
+    "# if using google\n",
+    "# export GCS_MODEL_URI=\"gs://CHANGEME\" # update with your value\n",
+    "# gsutil cp -r \"${LOCAL_MODEL_PATH}\" \"${GCS_MODEL_URI}\""
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

rasa/model_training.py

@@ -157,7 +157,6 @@ async def train(
     finetuning_epoch_fraction: float = 1.0,
     remote_storage: Optional[StorageType] = None,
     file_importer: Optional[TrainingDataImporter] = None,
-    keep_local_model_copy: bool = False,
 ) -> TrainingResult:
     """Trains a Rasa model (Core and NLU).
 
@@ -183,8 +182,6 @@ async def train(
             use for storing the model.
         file_importer: Instance of `TrainingDataImporter` to use for training.
             If it is not provided, a new instance will be created.
-        keep_local_model_copy: If `True` the model will be stored locally even if
-            remote storage is configured.
 
     Returns:
         An instance of `TrainingResult`.
@@ -266,7 +263,6 @@ async def train(
             finetuning_epoch_fraction=finetuning_epoch_fraction,
             dry_run=dry_run,
             remote_storage=remote_storage,
-            keep_local_model_copy=keep_local_model_copy,
             **(core_additional_arguments or {}),
             **(nlu_additional_arguments or {}),
         )
@@ -281,7 +277,6 @@ async def _train_graph(
     force_full_training: bool = False,
     dry_run: bool = False,
     remote_storage: Optional[StorageType] = None,
-    keep_local_model_copy: bool = False,
     **kwargs: Any,
 ) -> TrainingResult:
     if model_to_finetune:
@@ -322,10 +317,6 @@ async def _train_graph(
     rasa.engine.validation.validate_coexistance_routing_setup(
         domain, model_configuration, flows
     )
-    rasa.engine.validation.validate_model_group_configuration_setup()
-    rasa.engine.validation.validate_model_client_configuration_setup_during_training_time(
-        config
-    )
     rasa.engine.validation.validate_flow_component_dependencies(
         flows, model_configuration
     )
@@ -348,7 +339,7 @@ async def _train_graph(
             )
             return _dry_run_result(fingerprint_status, force_full_training)
 
-        model_name = determine_model_name(fixed_model_name, training_type)
+        model_name = _determine_model_name(fixed_model_name, training_type)
         full_model_path = Path(output_path, model_name)
 
         with telemetry.track_model_training(
@@ -363,8 +354,7 @@ async def _train_graph(
             )
             if remote_storage:
                 push_model_to_remote_storage(full_model_path, remote_storage)
-                if not keep_local_model_copy:
-                    full_model_path.unlink()
+                full_model_path.unlink()
                 structlogger.info(
                     "model_training.train.finished_training",
                     event_info=(
@@ -396,14 +386,9 @@ def _create_model_storage(
     return model_storage
 
 
-def generate_random_model_name() -> str:
-    time_format = "%Y%m%d-%H%M%S"
-    return f"{time.strftime(time_format)}-{randomname.get_name()}"
-
-
-def determine_model_name(
+def _determine_model_name(
     fixed_model_name: Optional[Text], training_type: TrainingType
-) -> str:
+) -> Text:
     if fixed_model_name:
         if not fixed_model_name.endswith(".tar.gz"):
             return f"{fixed_model_name}.tar.gz"
@@ -413,7 +398,8 @@ def determine_model_name(
     if training_type in [TrainingType.CORE, TrainingType.NLU]:
         prefix = f"{training_type.model_type}-"
 
-    return f"{prefix}{generate_random_model_name()}.tar.gz"
+    time_format = "%Y%m%d-%H%M%S"
+    return f"{prefix}{time.strftime(time_format)}-{randomname.get_name()}.tar.gz"
 
 
 async def train_core(
@@ -425,7 +411,6 @@ async def train_core(
     additional_arguments: Optional[Dict] = None,
     model_to_finetune: Optional[Text] = None,
     finetuning_epoch_fraction: float = 1.0,
-    keep_local_model_copy: bool = False,
 ) -> Optional[Text]:
     """Trains a Core model.
 
@@ -440,8 +425,6 @@ async def train_core(
             a directory in case the latest trained model should be used.
         finetuning_epoch_fraction: The fraction currently specified training epochs
             in the model configuration which should be used for finetuning.
-        keep_local_model_copy: If `True` the model will be stored locally even if
-            remote storage is configured.
 
     Returns:
         Path to the model archive.
@@ -497,7 +480,6 @@ async def train_core(
             model_to_finetune=model_to_finetune,
             fixed_model_name=fixed_model_name,
             finetuning_epoch_fraction=finetuning_epoch_fraction,
-            keep_local_model_copy=keep_local_model_copy,
             **(additional_arguments or {}),
         )
     ).model
@@ -513,7 +495,6 @@ async def train_nlu(
     domain: Optional[Union[Domain, Text]] = None,
     model_to_finetune: Optional[Text] = None,
     finetuning_epoch_fraction: float = 1.0,
-    keep_local_model_copy: bool = False,
 ) -> Optional[Text]:
     """Trains an NLU model.
 
@@ -531,8 +512,6 @@ async def train_nlu(
             a directory in case the latest trained model should be used.
         finetuning_epoch_fraction: The fraction currently specified training epochs
             in the model configuration which should be used for finetuning.
-        keep_local_model_copy: If `True` the model will be stored locally even if
-            remote storage is configured.
 
     Returns:
         Path to the model archive.
@@ -574,14 +553,13 @@ async def train_nlu(
             fixed_model_name=fixed_model_name,
             finetuning_epoch_fraction=finetuning_epoch_fraction,
             persist_nlu_training_data=persist_nlu_training_data,
-            keep_local_model_copy=keep_local_model_copy,
             **(additional_arguments or {}),
         )
     ).model
 
 
 def push_model_to_remote_storage(model_path: Path, remote_storage: StorageType) -> None:
-    """Push model to remote storage."""
+    """push model to remote storage"""
     from rasa.core.persistor import get_persistor
 
     persistor = get_persistor(remote_storage)

rasa/nlu/classifiers/diet_classifier.py

@@ -1,17 +1,18 @@
 from __future__ import annotations
-
 import copy
 import logging
 from collections import defaultdict
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Text, Tuple, Union, TypeVar, Type
+
+from rasa.exceptions import ModelNotFound
+from rasa.nlu.featurizers.featurizer import Featurizer
 
 import numpy as np
 import scipy.sparse
 import tensorflow as tf
 
-from rasa.exceptions import ModelNotFound
-from rasa.nlu.featurizers.featurizer import Featurizer
+from typing import Any, Dict, List, Optional, Text, Tuple, Union, TypeVar, Type
+
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
@@ -19,21 +20,18 @@ from rasa.engine.storage.storage import ModelStorage
 from rasa.nlu.extractors.extractor import EntityExtractorMixin
 from rasa.nlu.classifiers.classifier import IntentClassifier
 import rasa.shared.utils.io
+import rasa.utils.io as io_utils
 import rasa.nlu.utils.bilou_utils as bilou_utils
 from rasa.shared.constants import DIAGNOSTIC_DATA
 from rasa.nlu.extractors.extractor import EntityTagSpec
 from rasa.nlu.classifiers import LABEL_RANKING_LENGTH
 from rasa.utils import train_utils
 from rasa.utils.tensorflow import rasa_layers
-from rasa.utils.tensorflow.feature_array import (
-    FeatureArray,
-    serialize_nested_feature_arrays,
-    deserialize_nested_feature_arrays,
-)
 from rasa.utils.tensorflow.models import RasaModel, TransformerRasaModel
 from rasa.utils.tensorflow.model_data import (
     RasaModelData,
     FeatureSignature,
+    FeatureArray,
 )
 from rasa.nlu.constants import TOKENS_NAMES, DEFAULT_TRANSFORMER_SIZE
 from rasa.shared.nlu.constants import (
@@ -120,6 +118,7 @@ LABEL_SUB_KEY = IDS
 
 POSSIBLE_TAGS = [ENTITY_ATTRIBUTE_TYPE, ENTITY_ATTRIBUTE_ROLE, ENTITY_ATTRIBUTE_GROUP]
 
+
 DIETClassifierT = TypeVar("DIETClassifierT", bound="DIETClassifier")
 
 
@@ -1084,24 +1083,18 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
 
             self.model.save(str(tf_model_file))
 
-            # save data example
-            serialize_nested_feature_arrays(
-                self._data_example,
-                model_path / f"{file_name}.data_example.st",
-                model_path / f"{file_name}.data_example_metadata.json",
-            )
-            # save label data
-            serialize_nested_feature_arrays(
-                dict(self._label_data.data) if self._label_data is not None else {},
-                model_path / f"{file_name}.label_data.st",
-                model_path / f"{file_name}.label_data_metadata.json",
+            io_utils.pickle_dump(
+                model_path / f"{file_name}.data_example.pkl", self._data_example
             )
-
-            rasa.shared.utils.io.dump_obj_as_json_to_file(
-                model_path / f"{file_name}.sparse_feature_sizes.json",
+            io_utils.pickle_dump(
+                model_path / f"{file_name}.sparse_feature_sizes.pkl",
                 self._sparse_feature_sizes,
             )
-            rasa.shared.utils.io.dump_obj_as_json_to_file(
+            io_utils.pickle_dump(
+                model_path / f"{file_name}.label_data.pkl",
+                dict(self._label_data.data) if self._label_data is not None else {},
+            )
+            io_utils.json_pickle(
                 model_path / f"{file_name}.index_label_id_mapping.json",
                 self.index_label_id_mapping,
             )
@@ -1190,22 +1183,15 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
     ]:
         file_name = cls.__name__
 
-        # load data example
-        data_example = deserialize_nested_feature_arrays(
-            str(model_path / f"{file_name}.data_example.st"),
-            str(model_path / f"{file_name}.data_example_metadata.json"),
+        data_example = io_utils.pickle_load(
+            model_path / f"{file_name}.data_example.pkl"
         )
-        # load label data
-        loaded_label_data = deserialize_nested_feature_arrays(
-            str(model_path / f"{file_name}.label_data.st"),
-            str(model_path / f"{file_name}.label_data_metadata.json"),
-        )
-        label_data = RasaModelData(data=loaded_label_data)
-
-        sparse_feature_sizes = rasa.shared.utils.io.read_json_file(
-            model_path / f"{file_name}.sparse_feature_sizes.json"
+        label_data = io_utils.pickle_load(model_path / f"{file_name}.label_data.pkl")
+        label_data = RasaModelData(data=label_data)
+        sparse_feature_sizes = io_utils.pickle_load(
+            model_path / f"{file_name}.sparse_feature_sizes.pkl"
         )
-        index_label_id_mapping = rasa.shared.utils.io.read_json_file(
+        index_label_id_mapping = io_utils.json_unpickle(
             model_path / f"{file_name}.index_label_id_mapping.json"
         )
         entity_tag_specs = rasa.shared.utils.io.read_json_file(
@@ -1225,6 +1211,7 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
             for tag_spec in entity_tag_specs
         ]
 
+        # jsonpickle converts dictionary keys to strings
         index_label_id_mapping = {
             int(key): value for key, value in index_label_id_mapping.items()
         }