From fc6eab9edc868d2328936a18d19e1b37954cc60a Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Fri, 6 Jan 2023 17:05:47 +0000
Subject: [PATCH] [fix] new code complete answer and update readme for clarity

---
 model/supervised_finetuning/README.md         | 32 ++++++++++++++++++-
 .../custom_datasets/prompt_dialogue.py        |  7 ++--
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/model/supervised_finetuning/README.md b/model/supervised_finetuning/README.md
index 014afa95..ec337a3d 100644
--- a/model/supervised_finetuning/README.md
+++ b/model/supervised_finetuning/README.md
@@ -23,7 +23,37 @@ open-asisstant dataset are available it will be added here.
 
 ## Model
 
-TBD
+Normally you should be able to add new models in configs/config.yml
+
+```
+your-model-name:
+  learning_rate: 2e-6
+  model_name: <huggingface model name>
+  weight_decay: 0.01
+  max_length: 812
+  warmup_steps: 600
+  gradient_checkpointing: false
+  gradient_accumulation_steps: 5
+  per_device_train_batch_size: 4
+  per_device_eval_batch_size: 4
+```
+
+```
+python trainer.py --configs defaults your-model-name
+```
+
+However, if the model of your choice doesn't have pad_token, eos_token, sep_token, you have to update utils.py `get_tokenizer` to use the right token.
+
+
+## Deepspeed support
+
+You can edit the configs/zero_config.json and use any stage you wish. The current config uses zero-stage 3. For more details on how to setup the config checkout [this page](https://www.deepspeed.ai/tutorials/zero/)
+
+Once you are satisfy with your deepzero config, you can add --deepspeed flag at the end to trigger deepspeed
+
+```
+python trainer.py --configs defaults your-model-name --deepspeed
+```
 
 ## Results
 
diff --git a/model/supervised_finetuning/custom_datasets/prompt_dialogue.py b/model/supervised_finetuning/custom_datasets/prompt_dialogue.py
index 08a09113..17911141 100644
--- a/model/supervised_finetuning/custom_datasets/prompt_dialogue.py
+++ b/model/supervised_finetuning/custom_datasets/prompt_dialogue.py
@@ -30,17 +30,16 @@ class PromptGeneratedDataset(Dataset):
         answer = ""
         self.pairs = []
         with open(chat_dialogue, "r") as f:
-            corpus = f.read().split('<|endoftext|>')
+            corpus = f.read().split("<|endoftext|>")
             for dialogue in corpus:
                 dialogue = dialogue.strip()
-                if 'Rosey:' in dialogue:
-                    user, bot = dialogue.split('Rosey:', maxsplit=1)
+                if "Rosey:" in dialogue:
+                    user, bot = dialogue.split("Rosey:", maxsplit=1)
                     question = user.split(":", maxsplit=1)[1].strip()
                     answer = bot.strip()
                     if len(answer) and len(question):
                         self.pairs.append((question, answer))
 
-
         if len(question) > 0 and len(answer) > 0:
             self.pairs.append((question, answer))