Update README.md

README.md

@@ -13,16 +13,17 @@ Our overall explanation models along with ablations can be found in our [paper](
 |-|-|
 |[Flipped_11B](https://huggingface.co/seonghyeonye/flipped_11B)|11 billion|
 |[Flipped_3B](https://huggingface.co/seonghyeonye/flipped_3B)|3 billion|
-Here is how to use the model in PyTorch:
+Here is how to download the model in PyTorch:
+
 ```python
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-tokenizer = AutoTokenizer.from_pretrained("seonghyeonye/flipped_3B")
-model = AutoModelForSeq2SeqLM.from_pretrained("seonghyeonye/flipped_3B")
-inputs = tokenizer.encode("input: this is the best cast iron skillet you will ever buy\noutput: Positive", return_tensors="pt")
-outputs = model.generate(inputs)
-print(tokenizer.decode(outputs[0]))
+import torch
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+model = T5ForConditionalGeneration.from_pretrained("seonghyeonye/flipped_11B")
+tokenizer = T5Tokenizer.from_pretrained("seonghyeonye/flipped_11B")
 ```
-If you want to use another checkpoint, please replace the path in `AutoTokenizer` and `AutoModelForSeq2SeqLM`.
+If you want to use another checkpoint, please replace the path in `T5Tokenizer` and `T5ForConditionalGeneration`.
+We also provide a quick [Jupyter Notebook](https://github.com/seonghyeonye/Flipped-Learning/blob/master/flipped_inference.ipynb) where you can inference with our method.
 **Note: the model was trained with fp32 activations. As such, we highly discourage running inference with fp16.**
 
 # Training procedure