Update README.md

README.md

@@ -28,16 +28,17 @@ This information and more can also be found:
 
 # Use
 
-Install `transformers` **from [this source](https://github.com/swj0419/transformers_flexolmo)** and run:
+Install `transformers` with version `4.57.0` or newer and run:
 ```python
-from transformers import Olmoe2ForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-MODEL_NAME = "allenai/FlexOlmo-7x7B-1T"
-model = Olmoe2ForCausalLM.from_pretrained(MODEL_NAME).to(DEVICE)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+MODEL_NAME = "allenai/Flex-code-2x7B-1T"
+TOKENIZER_NAME = "allenai/dolma2-tokenizer"
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(DEVICE)
+tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
 inputs = tokenizer("Bitcoin is", return_tensors="pt")
 inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
 out = model.generate(**inputs, max_length=64)