Update README.md

README.md

@@ -11,8 +11,8 @@ library_name: transformers
 
 - **Homepage: In Progress** 
 - **Repository: https://github.com/kaistAI/Janus** 
-- **Paper:** 
-- **Point of Contact:seongyun@kaist.ac.kr**
+- **Paper: https://arxiv.org/abs/2405.17977** 
+- **Point of Contact: suehyunpark@kaist.ac.kr**
 
 # TL; DR
 ![image/png](https://cdn-uploads.huggingface.co/production/uploads/6550c4f27bbfce1878f5f280/vrQl8D8FV3vqUJYbPgsiG.png)
@@ -20,17 +20,17 @@ library_name: transformers
 Janus is a model trained using [Mistral-7B-v0.2](https://huggingface.co/mistral-community/Mistral-7B-v0.2) as its base model. Janus has been trained on [Multifaceted Collection](https://huggingface.co/datasets/kaist-ai/Multifaceted-Collection-SFT), a preference dataset containing 196k unique system messages for aligning LLMs to diverse human preferences. Janus not only excels at generating personalized responses that cater to various human preferences but is also adept at producing responses that are generally preferred for being helpful and harmless.
 
 # Model Details
-Janus-RM-7B is a reward model created by training Janus with Multifaceted-Collection-RM. Janus-RM-7B generates rewards when provided with various system messages and instructions, along with the personalized responses generated in accordance with these. This can be utilized to perform tasks such as PPO and best-of-n sampling.
+Janus-RM-7B is a reward model created by training Janus-7B (which is trained for only 1 epoch on the full 196k training instances) with [Multifaceted-Collection-RM](https://huggingface.co/datasets/kaist-ai/Multifaceted-Collection-RM) and a similar-sized mix of representative general helpfulness data: 72% of [HH-RLHF](https://huggingface.co/datasets/Anthropic/hh-rlhf), 14% of [OASST1 dataset preprocessed for reward modeling](https://huggingface.co/datasets/tasksource/oasst1_pairwise_rlhf_reward), and 14% of [WebGPT Comparisons](https://huggingface.co/datasets/openai/webgpt_comparisons). Janus-RM-7B predicts a scalar reward when provided with a concatenation of system message, instruction, chosen response, and rejected response. This can be utilized to perform as a scoring function for Best-of-N sampling or for preference tuning with proximal policy optimization (PPO).
 
 ## Model Description
 
 - **Model type:** Language model
 - **Language(s) (NLP):** English
 - **License:** Apache 2.0
-- **Related Models:** [Janus-66k-7B]() [Janus-DPO-7B](), [Janus-ORPO-7B](), [Janus-7B]()
-- **Training Datasets**: [Multifaceted-Collection-SFT](https://huggingface.co/datasets/kaist-ai/Multifaceted-Collection-SFT)
+- **Related Models:** [Janus-DPO-7B](https://huggingface.co/kaist-ai/janus-dpo-7b), [Janus-ORPO-7B](https://huggingface.co/kaist-ai/janus-orpo-7b), [Janus-7B](https://huggingface.co/kaist-ai/janus-7b)
+- **Training Datasets**: [Multifaceted-Collection-RM](https://huggingface.co/datasets/kaist-ai/Multifaceted-Collection-RM), [Anthropic/hh-rlhf](https://huggingface.co/datasets/Anthropic/hh-rlhf), [tasksource/oasst1_pairwise_rlhf_reward](https://huggingface.co/datasets/tasksource/oasst1_pairwise_rlhf_reward), [openai/webgpt_comparisons](https://huggingface.co/datasets/openai/webgpt_comparisons)
 - **Resources for more information:**
-  - [Research paper]()
+  - [Research paper](https://arxiv.org/abs/2405.17977)
   - [GitHub Repo](https://github.com/kaistAI/Janus)
 
 # Usage
@@ -79,23 +79,28 @@ print(decoded[0][len(input_str):])
 ```
 To train Janus and evaluate the responses it generates, please refer to the [GitHub Repo](https://github.com/kaistAI/Janus).
 Additionally, refer to the [Multifaceted Bench](https://huggingface.co/datasets/kaist-ai/Multifaceted-Bench), which evaluates how well LLM generates personalized responses.
+
 # Training Details
 ## Training hyperparameters
 
-The following hyperparameters were used during training:
-- learning_rate: 5e-06
-- train_batch_size: 2
+The following hyperparameters were used for training:
+- learning_rate: 9e-6
+- train_batch_size: 8
 - eval_batch_size: 2
 - seed: 42
 - distributed_type: multi-GPU
 - num_devices: 4
 - gradient_accumulation_steps: 4
-- total_train_batch_size: 32
+- total_train_batch_size: 128
 - total_eval_batch_size: 8
-- optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
+- optimizer: AdamW with betas=(0.9,0.95)
 - lr_scheduler_type: cosine
-- lr_scheduler_warmup_steps: 10
-- num_epochs: 4
+- lr_scheduler_warmup_steps: 3% of the maximum number of steps
+- num_epochs: 1
+- use_flash_attention_2: true
+- maximum_sequence_length: 2048
+- bf16: true
+- gradient_checkpointing: true
 
 ## Framework versions
 
@@ -103,6 +108,7 @@ The following hyperparameters were used during training:
 - Pytorch 2.2.2
 - Datasets 2.18.0
 - Tokenizers 0.15.0
+- DeepSpeed Zero-3
 
 # Citation