[feat] Adjust naming according to UrarTU

UKPLab · Jul 18, 2024 · 3c57580 · 3c57580
1 parent 0efabc8
commit 3c57580
Show file tree

Hide file tree

Showing 17 changed files with 71 additions and 87 deletions.
diff --git a/README.md b/README.md
@@ -79,12 +79,12 @@ aim up
 Let's get started with generating dialogues using the `llm-roleplay` action. The process is simple: just provide the name of the configuration file containing the action, followed by the action name itself. For the `llm-roleplay` action, we'll initiate it by using the Mistral 8x7B model as the inquirer. 🎇
 
 ```bash
-urartu launch --name=llm_roleplay action_config=dialogue_generator +action_config/task/model_inquirer=mixtral +action_config/task/model_responder=llama action_config.task.model_inquirer.api_token="YOUR_TOKEN"
+urartu launch --name=llm_roleplay action_config=dialogue_generator aim=aim slurm=slurm +action_config/task/model_inquirer=mixtral +action_config/task/model_responder=llama action_config.task.model_inquirer.api_token="YOUR_TOKEN"
 ```
 
-The `action_config` parameter specifies which configuration file to use to run the action. Afterward, we define the configuration file for the inquirer using the `model_inquirer` argument and set the configuration for the responder with the `model_responder` argument.
+The `aim` and `slurm` configs read the Aim and Slurm configurations from `aim` and `slurm` files which are located in `llm_roleplay/configs_{username}/aim/aim.yaml` and `llm_roleplay/configs_{username}/slurm/slurm.yaml` respectively. The `action_config` parameter specifies which configuration file to use to run the action. Afterward, we define the configuration file for the inquirer using the `model_inquirer` argument and set the configuration for the responder with the `model_responder` argument.
 
-To execute the command on a Slurm cluster, modify the `llm_roleplay/configs/action_config/dialogue_generator.yaml` file with the corresponding fields, and then use the same command to run the job. For more details on how to edit the configuration files, please refer to the upcoming sections.
+To execute the command on a Slurm cluster, modify the `llm_roleplay/configs_{username}/slurm/slurm.yaml` file with the corresponding fields, and then use the same command to run the job. For more details on how to edit the configuration files, please refer to the upcoming sections.
 
 > **Huggingface Authentication**
 > You might need to log in to HuggingFace to authenticate your use of Mistral 8x7B. To do this, use the `huggingface-cli` login command and provide your access token.

diff --git a/llm_roleplay/VERSION b/llm_roleplay/VERSION
@@ -1 +1 @@
-2.0.5
+2.0.6
diff --git a/llm_roleplay/configs/action_config/dialogue_generator.yaml b/llm_roleplay/configs/action_config/dialogue_generator.yaml
@@ -2,6 +2,11 @@
 action_name: dialogue_generator
 seed: 5
 
+hydra:
+  sweeper:
+    params:
+      seed: 73,321,8479
+
 action_config:
   workdir: "./"
   experiment_name: generate dialogues
@@ -15,7 +20,7 @@ action_config:
 
     dataset:
       type:
-        _target_: llm_roleplay.datasets.hf_datasets.HFDatasets
+        _target_: urartu.datasets.hf.dataset_from_dict.DatasetFromDict
       input_key: "instruction"
       data:
         instruction:
@@ -195,46 +200,3 @@ action_config:
         native_english:
           - "is"
           - "is not"
-
-  dataloader:
-    batch_size: 16
-    num_workers: 4
-    shuffle: false
-
-slurm:
-  # Whether or not to run the job on SLURM
-  use_slurm: false
-  # Name of the job on SLURM
-  name: "example"
-  # Comment of the job on SLURM
-  comment: "example job"
-  # Partition of SLURM on which to run the job. This is a required field if using SLURM.
-  partition: ""
-  account: ""
-  # Where the logs produced by the SLURM jobs will be output
-  log_folder: "./slurm_logs"
-  # Maximum number of hours / minutes needed by the job to complete. Above this limit, the job might be pre-empted.
-  time_hours: 1
-  time_minutes: 0
-  # Additional constraints on the hardware of the nodes to allocate (example 'volta' to select a volta GPU)
-  constraint: ""
-  # GB of RAM memory to allocate for each node
-  mem_gb: 100
-  # TCP port on which the workers will synchronize themselves with torch distributed
-  port_id: 40050
-  # Number of CPUs per GPUs to request on the cluster.
-  num_cpu_per_proc: 4
-  # Number of GPUs per node to request on the cluster.
-  num_gpu_per_node: 4
-  # Number of nodes to request on the cluster.
-  num_nodes: 1
-  # Number of processes per node to request on the cluster.
-  num_proc_per_node: 1
-  # Any other parameters for slurm (e.g. account, hint, distribution, etc.,) as dictated by submitit.
-  # Please see https://github.com/facebookincubator/submitit/issues/23#issuecomment-695217824.
-  additional_parameters: {}
-
-hydra:
-  sweeper:
-    params:
-      seed: 73,321,8479
diff --git a/llm_roleplay/configs/action_config/task/model_inquirer/falcon.yaml b/llm_roleplay/configs/action_config/task/model_inquirer/falcon.yaml
@@ -1,5 +1,5 @@
 type:
-  _target_: llm_roleplay.models.causal_lm_model.CausalLMModel
+  _target_: llm_roleplay.models.model_causal_language.ModelCausalLanguage
 name: "tiiuae/falcon-40b-instruct"
 cache_dir: ""
 dtype: torch.float16

diff --git a/llm_roleplay/configs/action_config/task/model_inquirer/gpt3.5.yaml b/llm_roleplay/configs/action_config/task/model_inquirer/gpt3.5.yaml
@@ -1,5 +1,5 @@
 type:
-  _target_: llm_roleplay.models.openai_model.OpenAIModel
+  _target_: llm_roleplay.models.model_openai.ModelOpenAI
 openai_api_type: "azure"
 openai_api_version: "2023-05-15"
 azure_openai_endpoint: null

diff --git a/llm_roleplay/configs/action_config/task/model_inquirer/gpt4.yaml b/llm_roleplay/configs/action_config/task/model_inquirer/gpt4.yaml
@@ -1,5 +1,5 @@
 type:
-  _target_: llm_roleplay.models.openai_model.OpenAIModel
+  _target_: llm_roleplay.models.model_openai.ModelOpenAI
 openai_api_type: "azure"
 openai_api_version: "2023-05-15"
 azure_openai_endpoint: null

diff --git a/llm_roleplay/configs/action_config/task/model_inquirer/llama.yaml b/llm_roleplay/configs/action_config/task/model_inquirer/llama.yaml
@@ -1,5 +1,5 @@
 type:
-  _target_: llm_roleplay.models.causal_lm_model.CausalLMModel
+  _target_: llm_roleplay.models.model_causal_language.ModelCausalLanguage
 name: "meta-llama/Llama-2-13b-chat-hf"
 cache_dir: ""
 dtype: torch.float16

diff --git a/llm_roleplay/configs/action_config/task/model_inquirer/mixtral.yaml b/llm_roleplay/configs/action_config/task/model_inquirer/mixtral.yaml
@@ -1,5 +1,5 @@
 type:
-  _target_: llm_roleplay.models.causal_lm_model.CausalLMModel
+  _target_: llm_roleplay.models.model_causal_language.ModelCausalLanguage
 name: "mistralai/Mixtral-8x7B-Instruct-v0.1"
 role: "inquirer"
 cache_dir: ""

diff --git a/llm_roleplay/configs/action_config/task/model_inquirer/vicuna.yaml b/llm_roleplay/configs/action_config/task/model_inquirer/vicuna.yaml
@@ -1,5 +1,5 @@
 type:
-  _target_: llm_roleplay.models.causal_lm_model.CausalLMModel
+  _target_: llm_roleplay.models.model_causal_language.ModelCausalLanguage
 name: "lmsys/vicuna-13b-v1.5-16k"
 cache_dir: ""
 dtype: torch.float16

diff --git a/llm_roleplay/configs/action_config/task/model_responder/llama.yaml b/llm_roleplay/configs/action_config/task/model_responder/llama.yaml
@@ -1,5 +1,5 @@
 type:
-  _target_: llm_roleplay.models.pipeline_model.PipelineModel
+  _target_: llm_roleplay.models.model_pipeline.ModelPipeline
 name: "models--llama-2-hf/13B-Chat"
 cache_dir: ""
 dtype: torch.float16

diff --git a/llm_roleplay/configs_tamoyan/aim/aim.yaml b/llm_roleplay/configs_tamoyan/aim/aim.yaml
@@ -0,0 +1,2 @@
+repo: aim://0.0.0.0:53800
+log_system_params: true
diff --git a/llm_roleplay/configs_tamoyan/slurm/slurm.yaml b/llm_roleplay/configs_tamoyan/slurm/slurm.yaml
@@ -0,0 +1,32 @@
+
+# Whether or not to run the job on SLURM
+use_slurm: false
+# Name of the job on SLURM
+name: "example"
+# Comment of the job on SLURM
+comment: "example job"
+# Partition of SLURM on which to run the job. This is a required field if using SLURM.
+partition: ""
+account: ""
+# Where the logs produced by the SLURM jobs will be output
+log_folder: "./slurm_logs"
+# Maximum number of hours / minutes needed by the job to complete. Above this limit, the job might be pre-empted.
+time_hours: 1
+time_minutes: 0
+# Additional constraints on the hardware of the nodes to allocate (example 'volta' to select a volta GPU)
+constraint: ""
+# GB of RAM memory to allocate for each node
+mem_gb: 100
+# TCP port on which the workers will synchronize themselves with torch distributed
+port_id: 40050
+# Number of CPUs per GPUs to request on the cluster.
+num_cpu_per_proc: 4
+# Number of GPUs per node to request on the cluster.
+num_gpu_per_node: 4
+# Number of nodes to request on the cluster.
+num_nodes: 1
+# Number of processes per node to request on the cluster.
+num_proc_per_node: 1
+# Any other parameters for slurm (e.g. account, hint, distribution, etc.,) as dictated by submitit.
+# Please see https://github.com/facebookincubator/submitit/issues/23#issuecomment-695217824.
+additional_parameters: {}
diff --git a/llm_roleplay/datasets/hf_datasets.py b/llm_roleplay/datasets/hf_datasets.py
diff --git a/llm_roleplay/models/causal_lm_model.py → llm_roleplay/models/model_causal_language.py b/llm_roleplay/models/causal_lm_model.py → llm_roleplay/models/model_causal_language.py
@@ -8,7 +8,7 @@
 from llm_roleplay.common.model import Model
 
 
-class CausalLMModel(Model):
+class ModelCausalLanguage(Model):
     SELF_REPLY_TOKENS = {
         "llama": "[INST",
         "vicuna": "### Human:",
@@ -129,7 +129,7 @@ def generate(self, prompt: str, generate_cfg):
         turn_response = output_o.replace(model_prompt_o, "", 1)
 
         # ----------------------------------- prevent potential self-reply -----------------------------------
-        for self_reply_token in CausalLMModel.SELF_REPLY_TOKENS.values():
+        for self_reply_token in ModelCausalLanguage.SELF_REPLY_TOKENS.values():
             if self_reply_token in turn_response:
                 turn_response = turn_response.split(self_reply_token)[0]
                 self.aim_run["num_self_replies"] += 1

diff --git a/llm_roleplay/models/openai_model.py → llm_roleplay/models/model_openai.py b/llm_roleplay/models/openai_model.py → llm_roleplay/models/model_openai.py
@@ -8,7 +8,7 @@
 from llm_roleplay.common.model import Model
 
 
-class OpenAIModel(Model):
+class ModelOpenAI(Model):
     def __init__(self, cfg, role) -> None:
         super().__init__(cfg, role)
 

diff --git a/llm_roleplay/models/pipeline_model.py → llm_roleplay/models/model_pipeline.py b/llm_roleplay/models/pipeline_model.py → llm_roleplay/models/model_pipeline.py
@@ -5,7 +5,7 @@
 from llm_roleplay.common.model import Model
 
 
-class PipelineModel(Model):
+class ModelPipeline(Model):
     SELF_REPLY_TOKENS = {
         "llama": "[INST",
         "vicuna": "### Human:",
@@ -122,7 +122,7 @@ def generate(self, prompt: str, generate_cfg):
         turn_response = output_o.replace(model_prompt_o, "", 1)
 
         # ----------------------------------- prevent potential self-reply -----------------------------------
-        for self_reply_token in PipelineModel.SELF_REPLY_TOKENS.values():
+        for self_reply_token in ModelPipeline.SELF_REPLY_TOKENS.values():
             if self_reply_token in turn_response:
                 turn_response = turn_response.split(self_reply_token)[0]
                 self.aim_run["num_self_replies"] += 1

diff --git a/tests/test_roleplay.py b/tests/test_roleplay.py
@@ -24,7 +24,7 @@ def setUp(self):
                     "task": {
                         "num_turns": 2,
                         "model_inquirer": {
-                            "type": {"_target_": "llm_roleplay.models.causal_lm_model.CausalLMModel"},
+                            "type": {"_target_": "llm_roleplay.models.model_causal_language.ModelCausalLanguage"},
                             "name": "mistralai/Mixtral-8x7B-Instruct-v0.1",
                             "cache_dir": "",
                             "dtype": "torch.float16",
@@ -48,7 +48,7 @@ def setUp(self):
                             "idx_of_possible_prompt": 0,
                         },
                         "model_responder": {
-                            "type": {"_target_": "llm_roleplay.models.pipeline_model.PipelineModel"},
+                            "type": {"_target_": "llm_roleplay.models.model_pipeline.ModelPipeline"},
                             "name": "models--llama-2-hf/13B-Chat",
                             "cache_dir": "",
                             "dtype": "torch.float16",
@@ -63,7 +63,7 @@ def setUp(self):
                             },
                         },
                         "dataset": {
-                            "type": {"_target_": "llm_roleplay.datasets.hf_datasets.HFDatasets"},
+                            "type": {"_target_": "llm_roleplay.datasets.hf.dataset_from_dict.DatasetFromDict"},
                             "input_key": "instruction",
                             "data": {
                                 "instruction": [
@@ -122,11 +122,11 @@ def test_tracking_calls(self):
         self.assertIsNotNone(text_record, "No text records found in AIM run tracking")
 
     @patch("llm_roleplay.models.openai_model.OpenAIModel.generate")
-    @patch("llm_roleplay.models.pipeline_model.PipelineModel.generate")
-    @patch("llm_roleplay.models.causal_lm_model.CausalLMModel.generate")
+    @patch("llm_roleplay.models.model_pipeline.ModelPipeline.generate")
+    @patch("llm_roleplay.models.model_causal_language.ModelCausalLanguage.generate")
     @patch("llm_roleplay.models.openai_model.OpenAIModel.model")
-    @patch("llm_roleplay.models.pipeline_model.PipelineModel.model")
-    @patch("llm_roleplay.models.causal_lm_model.CausalLMModel.model")
+    @patch("llm_roleplay.models.model_pipeline.ModelPipeline.model")
+    @patch("llm_roleplay.models.model_causal_language.ModelCausalLanguage.model")
     def test_initialization(
         self,
         mock_model_clm,
@@ -183,11 +183,11 @@ def test_initialization(
 
 
     @patch("llm_roleplay.models.openai_model.OpenAIModel.generate")
-    @patch("llm_roleplay.models.pipeline_model.PipelineModel.generate")
-    @patch("llm_roleplay.models.causal_lm_model.CausalLMModel.generate")
+    @patch("llm_roleplay.models.model_pipeline.ModelPipeline.generate")
+    @patch("llm_roleplay.models.model_causal_language.ModelCausalLanguage.generate")
     @patch("llm_roleplay.models.openai_model.OpenAIModel.model")
-    @patch("llm_roleplay.models.pipeline_model.PipelineModel.model")
-    @patch("llm_roleplay.models.causal_lm_model.CausalLMModel.model")
+    @patch("llm_roleplay.models.model_pipeline.ModelPipeline.model")
+    @patch("llm_roleplay.models.model_causal_language.ModelCausalLanguage.model")
     @patch("torch.cuda.empty_cache")
     def test_resource_management(
         self,
@@ -224,11 +224,11 @@ def test_resource_management(
         mock_empty_cache.assert_called()
 
     @patch("llm_roleplay.models.openai_model.OpenAIModel.generate")
-    @patch("llm_roleplay.models.pipeline_model.PipelineModel.generate")
-    @patch("llm_roleplay.models.causal_lm_model.CausalLMModel.generate")
+    @patch("llm_roleplay.models.model_pipeline.ModelPipeline.generate")
+    @patch("llm_roleplay.models.model_causal_language.ModelCausalLanguage.generate")
     @patch("llm_roleplay.models.openai_model.OpenAIModel.model")
-    @patch("llm_roleplay.models.pipeline_model.PipelineModel.model")
-    @patch("llm_roleplay.models.causal_lm_model.CausalLMModel.model")
+    @patch("llm_roleplay.models.model_pipeline.ModelPipeline.model")
+    @patch("llm_roleplay.models.model_causal_language.ModelCausalLanguage.model")
     def test_dialogue_generation(
         self,
         mock_model_clm,