resolve comment on config model path.

CambioML · Sep 10, 2023 · 6896163 · 6896163
1 parent c48ae85
commit 6896163
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/example/rlhf/demo_rw_finetuning.py b/example/rlhf/demo_rw_finetuning.py
@@ -26,7 +26,6 @@
 
 # run reward model finetuning
 # config = pykoi.RLHFConfig(dataset_type="local_db")
-config = pykoi.RLHFConfig()
-config.base_model_path = "databricks/dolly-v2-3b"
+config = pykoi.RLHFConfig(reward_model_path = "databricks/dolly-v2-3b")
 rlhf_step2_rft = pykoi.RewardFinetuning(config)
 rlhf_step2_rft.train_and_save("./models/rlhf_step2_rw")
diff --git a/example/rlhf/supervised_finetuning_demo.py b/example/rlhf/supervised_finetuning_demo.py
@@ -25,7 +25,6 @@
 print("My local database has {} samples in total".format(my_data_pd.shape[0]))
 
 # run supervised finetuning
-config = pykoi.RLHFConfig(base_model_path="elinas/llama-7b-hf-transformers-4.29", dataset_type="local_db")
-config.base_model_path = "databricks/dolly-v2-3b"
+config = pykoi.RLHFConfig(base_model_path="databricks/dolly-v2-3b", dataset_type="local_db")
 rlhf_step1_sft = pykoi.SupervisedFinetuning(config)
 rlhf_step1_sft.train_and_save("./models/rlhf_step1_sft")