Update README.md

README.md

@@ -31,11 +31,11 @@ It is important to note that the primary intended use case of this model is to c
 # pip install -q transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-checkpoint = "HuggingFaceFW/ablation-model-fineweb-edu"
+model = "HuggingFaceFW/ablation-model-fineweb-edu"
 device = "cuda" # for GPU usage or "cpu" for CPU usage
 
-tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
+tokenizer = AutoTokenizer.from_pretrained(model)
+model = AutoModelForCausalLM.from_pretrained(model).to(device)
 
 inputs = tokenizer.encode("Machine Learning is", return_tensors="pt").to(device)
 outputs = model.generate(inputs)
@@ -48,12 +48,12 @@ We are releasing intermediate checkpoints for this model at intervals of every 1
 
 You can load a specific model revision with `transformers` using the argument `revision`:
 ```python
-model = AutoModelForCausalLM.from_pretrained(checkpoint, revision="step-001000-2BT")
+model = AutoModelForCausalLM.from_pretrained("HuggingFaceFW/ablation-model-fineweb-edu", revision="step-001000-2BT")
 ```
 You can access all the revisions for the models via the following code:
 ```python
 from huggingface_hub import list_repo_refs
-out = list_repo_refs(checkpoint)
+out = list_repo_refs("HuggingFaceFW/ablation-model-fineweb-edu")
 branches = [b.name for b in out.branches]
 ```