Update README.md

README.md

@@ -29,9 +29,9 @@ import torch
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load different ckpts via passing e.g. `revision=step10000-tokens41B`
-model = OlmoeForCausalLM.from_pretrained("OLMoE/OLMoE-1B-7B-Instruct").to(DEVICE)
-tokenizer = AutoTokenizer.from_pretrained("OLMoE/OLMoE-1B-7B-Instruct")
+# Load different ckpts via passing e.g. `revision=kto`
+model = OlmoeForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-Instruct").to(DEVICE)
+tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-Instruct")
 message = [{"role": "user", "content": "Explain to me like I'm five what is Bitcoin."}]
 inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
 out = model.generate(**inputs, max_length=64)
@@ -39,17 +39,11 @@ print(tokenizer.decode(out[0]))
 # > # Bitcoin is a digital currency that is created and held electronically. No one controls it. Bitcoins aren’t printed, like dollars or euros – they’re produced by people and businesses running computers all around the world, using software that solves mathematical
 ```
 
-You can list all revisions/branches by installing `huggingface-hub` & running:
-```python
-from huggingface_hub import list_repo_refs
-out = list_repo_refs("OLMoE/OLMoE-1B-7B-0824")
-branches = [b.name for b in out.branches]
-```
-
-Important branches:
+Branches:
 - `main`: Preference tuned via DPO model of https://hf.co/OLMoE/OLMoE-1B-7B-0824-SFT (`main` branch)
-- `no-load-balancing`: Ablation without load balancing loss during DPO starting from the `no-load-balancing` branch of https://hf.co/OLMoE/OLMoE-1B-7B-0824-SFT
+- `load-balancing`: Ablation with load balancing loss during DPO starting from the `load-balancing` branch of https://hf.co/OLMoE/OLMoE-1B-7B-0824-SFT
 - `non-annealed`: Ablation starting from the `non-annealed` branch of https://hf.co/OLMoE/OLMoE-1B-7B-0824-SFT which is an SFT of the pretraining checkpoint prior to annealing (branch `step1200000-tokens5033B` of https://hf.co/OLMoE/OLMoE-1B-7B-0824)
+- `kto`: Ablation using KTO instead of DPO. This branch is the checkpoint after 5,000 steps with the RMS optimizer. The other `kto*` branches correspond to the other checkpoints mentioned in the paper.
 
 # Citation