Update README.md
Browse files
README.md
CHANGED
|
@@ -35,6 +35,25 @@ DCLM-1B-v0 is a 1.4 billion parameter language model trained on the DCLM-Baselin
|
|
| 35 |
- **Dataset:** https://huggingface.co/datasets/mlfoundations/dclm-baseline-1.0
|
| 36 |
- **Paper:** [DataComp-LM: In search of the next generation of training sets for language models](https://arxiv.org/abs/2406.11794)
|
| 37 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 38 |
|
| 39 |
### Training Details
|
| 40 |
|
|
|
|
| 35 |
- **Dataset:** https://huggingface.co/datasets/mlfoundations/dclm-baseline-1.0
|
| 36 |
- **Paper:** [DataComp-LM: In search of the next generation of training sets for language models](https://arxiv.org/abs/2406.11794)
|
| 37 |
|
| 38 |
+
## Quickstart
|
| 39 |
+
First install open_lm
|
| 40 |
+
```
|
| 41 |
+
pip install git+https://github.com/mlfoundations/open_lm.git
|
| 42 |
+
```
|
| 43 |
+
|
| 44 |
+
Then you can load the model using HF's Auto classes as follows:
|
| 45 |
+
```python
|
| 46 |
+
from open_lm.hf import *
|
| 47 |
+
from transformers import AutoTokenizer, AutoModelForCausalLM
|
| 48 |
+
tokenizer = AutoTokenizer.from_pretrained("TRI-ML/DCLM-1B-v0")
|
| 49 |
+
model = AutoModelForCausalLM.from_pretrained("TRI-ML/DCLM-1B-v0")
|
| 50 |
+
|
| 51 |
+
inputs = tokenizer(["Machine learning is"], return_tensors="pt")
|
| 52 |
+
gen_kwargs = {"max_new_tokens": 50, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.1}
|
| 53 |
+
output = model.generate(inputs['input_ids'], **gen_kwargs)
|
| 54 |
+
output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
|
| 55 |
+
print(output)
|
| 56 |
+
```
|
| 57 |
|
| 58 |
### Training Details
|
| 59 |
|