Update README.md

README.md

@@ -4,12 +4,10 @@ datasets:
 - roneneldan/TinyStories
 ---
 Model trained on the TinyStories Dataset, replicating https://arxiv.org/abs/2305.07759
-
 Based on GPT-Neo architecture.
 
-
 hyperparams used to train this model:
-
+```
         "batch_size": 32,
         "block_size": 256,
         "lr": 5e-4,
@@ -24,17 +22,17 @@ hyperparams used to train this model:
         "vocab_size": 50257, 
         "warmup_tokens": 10000,
         "gradient_accumulation_steps": 8
-
+```
 ------ EXAMPLE USAGE ---
-
+```py
   !pip install --quiet transformers 
   from transformers import AutoModelForCausalLM, AutoTokenizer
   model20 = AutoModelForCausalLM.from_pretrained('AnirudhRajagopalan1201/tinystories-custom-20M')
   
   tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
   prompt = "Lily likes cats and dogs. She asked her mom for a dog and her mom said no, so instead she asked"
-  input_ids = tokenizer.encode(prompt, return_tensors="pt")
-  print("--------------20M parameter--------------------------------")
+  input_ids = tokenizer.encode(prompt, return_tensors="pt") 
   output = model20.generate(input_ids, temperature=0.2, max_length = 100, do_sample=True)
   output_text = tokenizer.decode(output[0], skip_special_tokens=True)
-  print(output_text)
+  print(output_text)
+```