diff --git a/lib/style-transfer/bash/experiment/rb_gen/az/llama3.2-3b.sh b/lib/style-transfer/bash/experiment/rb_gen/az/llama3.2-3b.sh
new file mode 100644
index 0000000..f7e39c2
--- /dev/null
+++ b/lib/style-transfer/bash/experiment/rb_gen/az/llama3.2-3b.sh
@@ -0,0 +1,8 @@
+python style_transfer/run_rb_gen.py model.name=meta-llama/Llama-3.2-3B-Instruct \
+            model.peft_config.target_modules='["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]' \
+            dataset.name=bio-datasets/mimic_style_transfer \
+            max_steps=5 \
+            dataset.num_generated_samples=500 \
+            score.model.model_name_or_path=FremyCompany/BioLORD-2023-C \
+            dataset.sft_ratio=0.06 \
+            dataset.gen_ratio=0.7
diff --git a/lib/style-transfer/bash/experiment/rb_gen/az/test-azure.sh b/lib/style-transfer/bash/experiment/rb_gen/az/test-azure.sh
new file mode 100644
index 0000000..2cb6cf2
--- /dev/null
+++ b/lib/style-transfer/bash/experiment/rb_gen/az/test-azure.sh
@@ -0,0 +1 @@
+python style_transfer/run_rb_gen.py
diff --git a/lib/style-transfer/bash/experiment/rb_gen/jz-slurm/llama3b.sh b/lib/style-transfer/bash/experiment/rb_gen/jz-slurm/llama3b.sh
new file mode 100644
index 0000000..fea8a67
--- /dev/null
+++ b/lib/style-transfer/bash/experiment/rb_gen/jz-slurm/llama3b.sh
@@ -0,0 +1,6 @@
+sbatch bash/jz-slurm/submit-a100.sh \
+            model.name=meta-llama/Llama-3.2-3B-Instruct \
+            model.lora.target_modules='["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]' \
+            dataset.name=bio-datasets/mimic_style_transfer \
+            max_steps=3 \
+            score.model_name_or_path=FremyCompany/BioLORD-2023-C
diff --git a/lib/style-transfer/bash/jz-slurm/singularity-conv.sh b/lib/style-transfer/bash/experiment/rb_gen/jz-slurm/singularity-conv.sh
similarity index 92%
rename from lib/style-transfer/bash/jz-slurm/singularity-conv.sh
rename to lib/style-transfer/bash/experiment/rb_gen/jz-slurm/singularity-conv.sh
index 594e324..6f7bcaa 100644
--- a/lib/style-transfer/bash/jz-slurm/singularity-conv.sh
+++ b/lib/style-transfer/bash/experiment/rb_gen/jz-slurm/singularity-conv.sh
@@ -4,7 +4,7 @@
 #SBATCH --cpus-per-task=10
 #SBATCH --output=log/conversion%j.out
 #SBATCH --partition=prepost
-#SBATCH --account=oha@cpu
+#SBATCH --account=oha@v100
 
 module purge
 module load singularity
diff --git a/lib/style-transfer/bash/experiment/rb_gen/jz-slurm/submit-a100.sh b/lib/style-transfer/bash/experiment/rb_gen/jz-slurm/submit-a100.sh
new file mode 100644
index 0000000..fca080a
--- /dev/null
+++ b/lib/style-transfer/bash/experiment/rb_gen/jz-slurm/submit-a100.sh
@@ -0,0 +1,21 @@
+#!/bin/bash
+#SBATCH --job-name=gpu-job
+#SBATCH --nodes=1
+#SBATCH --gres=gpu:1
+#SBATCH --cpus-per-task=10
+#SBATCH --output=log/test.out
+#SBATCH --partition=gpu_p5
+#SBATCH --account=oha@a100
+#SBATCH -C a100
+#SBATCH --qos=qos_gpu_a100-t3
+
+module purge
+module load singularity
+
+SCRIPT_PATH="style_transfer/run_rb_gen.py"
+PYTHON="/opt/pysetup/.venv/bin/python"
+LIB_PATH="$WORK/open-nlp/lib/style-transfer/"
+singularity exec --bind $WORK/open-nlp,$HF_HOME,$WANDB_CACHE_DIR,$WANDB_DIR \
+                 --env PYTHONPATH=$PYTHONPATH:$LIB_PATH,HF_HUB_OFFLINE=True \
+                 --nv $SINGULARITY_ALLOWED_DIR/style-transfer.sif \
+                 bash -c "cd $LIB_PATH && $PYTHON $SCRIPT_PATH $1"
diff --git a/lib/style-transfer/bash/jz-slurm/test-pipeline.sh b/lib/style-transfer/bash/experiment/rb_gen/jz-slurm/test-pipeline.sh
similarity index 95%
rename from lib/style-transfer/bash/jz-slurm/test-pipeline.sh
rename to lib/style-transfer/bash/experiment/rb_gen/jz-slurm/test-pipeline.sh
index 2965acd..5d854cb 100644
--- a/lib/style-transfer/bash/jz-slurm/test-pipeline.sh
+++ b/lib/style-transfer/bash/experiment/rb_gen/jz-slurm/test-pipeline.sh
@@ -7,7 +7,7 @@
 #SBATCH --partition=gpu_p5
 #SBATCH --account=oha@a100
 #SBATCH -C a100
-#SBATCH --qos=qos_gpu-dev
+#SBATCH --qos=qos_gpu_a100-dev
 
 module purge
 module load singularity
diff --git a/lib/style-transfer/bash/experiment/rb_gen/test-azure.sh b/lib/style-transfer/bash/experiment/rb_gen/test-azure.sh
deleted file mode 100644
index 8c4e2ba..0000000
--- a/lib/style-transfer/bash/experiment/rb_gen/test-azure.sh
+++ /dev/null
@@ -1 +0,0 @@
-python style_transfer/rb_gen.py
diff --git a/lib/style-transfer/configs/rb_gen/dpo/default.yaml b/lib/style-transfer/configs/rb_gen/dpo/default.yaml
index 5d584ad..4f599aa 100644
--- a/lib/style-transfer/configs/rb_gen/dpo/default.yaml
+++ b/lib/style-transfer/configs/rb_gen/dpo/default.yaml
@@ -19,9 +19,8 @@ training_args:
   report_to: "none"
 
 beta: 0.1
-threshold: 0.5
 checkpoint: null
 dataset: null
 max_length: 1024
 max_prompt_length: 512
-percentile: 60
+percentile: 50
diff --git a/lib/style-transfer/configs/rb_gen/train.yaml b/lib/style-transfer/configs/rb_gen/train.yaml
index 3341069..570ce11 100644
--- a/lib/style-transfer/configs/rb_gen/train.yaml
+++ b/lib/style-transfer/configs/rb_gen/train.yaml
@@ -33,6 +33,7 @@ dataset:
   name: bio-datasets/mimic_style_transfer_tiny
   sft_ratio: 0.5
   gen_ratio: 0.5
+  num_generated_samples: 2
 
 seed: 666
 max_steps: 1
diff --git a/lib/style-transfer/pyproject.toml b/lib/style-transfer/pyproject.toml
index dcbf793..e49ce40 100644
--- a/lib/style-transfer/pyproject.toml
+++ b/lib/style-transfer/pyproject.toml
@@ -64,7 +64,7 @@ rich = "^13.5.2"
 sentence-transformers = "^2.3.1"
 torch = "^2.3.0"
 tqdm = "^4.66.1"
-transformers = "4.42.3"
+transformers = "4.45.1"
 trl = "^0.9.6"
 vllm = "^0.5.0.post1"
 wandb = "^0.16.3"
diff --git a/lib/style-transfer/style_transfer/rb_gen/steps/generate.py b/lib/style-transfer/style_transfer/rb_gen/steps/generate.py
index 472eb07..0aa9623 100644
--- a/lib/style-transfer/style_transfer/rb_gen/steps/generate.py
+++ b/lib/style-transfer/style_transfer/rb_gen/steps/generate.py
@@ -27,7 +27,6 @@ def generate(
     best_model_path: str,
     tokenizer: PreTrainedTokenizerBase,
     gen_dataset: Dataset,
-    test_dataset: Dataset,
 ) -> Dataset:
     """Generate the synthetic candidates for the dataset.
     To improve the speed of the prediction we use VLLM pipeline.
@@ -40,7 +39,6 @@ def generate(
         best_model_path: The path to the best model.
         tokenizer: The tokenizer to use.
         gen_dataset: The dataset to generate the synthetic candidates for.
-        test_dataset: The dataset to generate the synthetic candidates for.
 
     Returns:
         The generated dataset.
@@ -56,8 +54,8 @@ def generate(
     tokenizer.save_pretrained("models/merged/")
     del model
     del tokenizer
-    gc.collect()
-    torch.cuda.empty_cache()
+    # gc.collect()
+    # torch.cuda.empty_cache()
     logging.info("🫧 Building VLLM Pipeline ...")
     llm = LLM(
         model="models/merged/",
@@ -66,24 +64,15 @@ def generate(
 
     logging.info("🎉 And it's done!")
 
+    shuffled_gen_dataset = gen_dataset.shuffle(seed=cfg.seed)
+    subset_gen_dataset = shuffled_gen_dataset.select(range(cfg.dataset.num_generated_samples))
     gen_dataloader = torch.utils.data.DataLoader(
-        gen_dataset,
-        batch_size=cfg.gen.batch_size,
-    )
-    test_dataloader = torch.utils.data.DataLoader(
-        test_dataset,
+        subset_gen_dataset if step != 0 else gen_dataset,
         batch_size=cfg.gen.batch_size,
     )
     gen_pred_dataset = batch_generate(
         cfg, step, gen_dataloader, llm, f"{wandb.config['state']}/gen_dataset"
     )
-    _ = batch_generate(
-        cfg,
-        step,
-        test_dataloader,
-        llm,
-        f"{wandb.config['state']}/test_dataset",
-    )
 
     wandb.log_artifact(CACHE_PATH, type="data")
     destroy_model_parallel()
diff --git a/lib/style-transfer/style_transfer/run_rb_gen.py b/lib/style-transfer/style_transfer/run_rb_gen.py
index f9965ca..73137d3 100644
--- a/lib/style-transfer/style_transfer/run_rb_gen.py
+++ b/lib/style-transfer/style_transfer/run_rb_gen.py
@@ -31,7 +31,7 @@ def main(cfg: DictConfig):
     """
     wandb.init(
         project="style-transfer",
-        id=f"{cfg.model.name.replace('/','-')}-{wandb.util.generate_id()}",
+        id=f"{cfg.model.name.replace('/', '-')}-{wandb.util.generate_id()}",
         resume="allow",
         settings=wandb.Settings(code_dir="."),
     )
@@ -76,7 +76,6 @@ def main(cfg: DictConfig):
             current_model_path,
             tokenizer,
             gen_dataset,
-            test_dataset,
         )
         score_dataset = score(
             cfg,
@@ -95,7 +94,6 @@ def main(cfg: DictConfig):
         current_model_path,
         tokenizer,
         gen_dataset,
-        test_dataset,
     )
     score(
         cfg,