Update README.md

README.md

@@ -42,8 +42,9 @@ In particular, we focus on four revisions of the 7B models:
 
 To load a specific model revision with HuggingFace, simply add the argument `revision`:
 ```bash
-import hf_olmo # pip install ai2-olmo
-olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B-Twin-2T", revision="step1000-tokens4B")
+from hf_olmo import OLMoForCausalLM  # pip install ai2-olmo
+
+olmo = OLMoForCausalLM.from_pretrained("allenai/OLMo-7B-Twin-2T", revision="step1000-tokens4B")
 ```
 
 All revisions/branches are listed in the file `revisions.txt`. 
@@ -93,11 +94,10 @@ pip install ai2-olmo
 ```
 Now, proceed as usual with HuggingFace:
 ```python
-import hf_olmo
+from hf_olmo import OLMoForCausalLM, OLMoTokenizerFast
 
-from transformers import AutoModelForCausalLM, AutoTokenizer
-olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B")
-tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-7B")
+olmo = OLMoForCausalLM.from_pretrained("allenai/OLMo-7B-Twin-2T")
+tokenizer = OLMoTokenizerFast.from_pretrained("allenai/OLMo-7B-Twin-2T")
 message = ["Language modeling is "]
 inputs = tokenizer(message, return_tensors='pt', return_token_type_ids=False)
 # optional verifying cuda
@@ -107,17 +107,8 @@ response = olmo.generate(**inputs, max_new_tokens=100, do_sample=True, top_k=50,
 print(tokenizer.batch_decode(response, skip_special_tokens=True)[0])
 >> 'Language modeling is the first step to build natural language generation...'
 ```
-Alternatively, with the pipeline abstraction:
-```python
-import hf_olmo
-
-from transformers import pipeline
-olmo_pipe = pipeline("text-generation", model="allenai/OLMo-7B")
-print(olmo_pipe("Language modeling is "))
->> 'Language modeling is a branch of natural language processing that aims to...'
-```
 
-Or, you can make this slightly faster by quantizing the model, e.g. `AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B", torch_dtype=torch.float16, load_in_8bit=True)` (requires `bitsandbytes`).
+You can make this slightly faster by quantizing the model, e.g. `AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B", torch_dtype=torch.float16, load_in_8bit=True)` (requires `bitsandbytes`).
 The quantized model is more sensitive to typing / cuda, so it is recommended to pass the inputs as `inputs.input_ids.to('cuda')` to avoid potential issues.
 
 Note, you may see the following error if `ai2-olmo` is not installed correctly, which is caused by internal Python check naming. We'll update the code soon to make this error clearer.