Update README.md

README.md

@@ -1,3 +1,4 @@
+```python
 from datasets import load_dataset
 
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -30,3 +31,4 @@ for prompt in val_prompts:
     count += 1
     if count == 10:
         break
+```