Update README.md

README.md

@@ -1,3 +1,19 @@
+---
+language: 
+  - en
+tags:
+  - gpt2
+license: apache-2.0
+datasets:
+  - wikitext
+  - openwebtext
+  - cc-stories
+metrics:
+  - type: wikitext
+    value: 19.31
+    name: WikiText-103
+---
+
 <!---
 # ##############################################################################################
 # 
@@ -48,29 +64,33 @@ model = GPT2LMHeadModel.from_pretrained("robowaifudev/megatron-gpt2-345m")
 
 if torch.cuda.is_available():
     device = torch.device("cuda")
-    model.to(device)
     model.half()
 else:
     device = torch.device("cpu")
+model.to(device)
 model.eval()
 
 # Generate
-text = "Hello world!"
-input_ids = tokenizer.encode(text, return_tensors="pt")
+prompt = "It was a bright cold day in April, and the clocks were striking thirteen. Winston Smith,"
+input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
 output = model.generate(
     input_ids=input_ids,
-    max_length=len(input_ids) + 32,
+    max_length=len(input_ids) + 128,
     do_sample=True,
     top_k=64,
     top_p=0.9,
     temperature=0.8,
-    num_return_sequences=1
+    num_return_sequences=2,
+    repetition_penalty=1.025
 )
 
 # Output the text.
+print("Prompt:", prompt)
+print("*" * 3)
 for i, sentence in enumerate(output):
     text = tokenizer.decode(sentence, clean_up_tokenization_spaces=True)
     print(f"{i}:", text)
+    print("*" * 3)
 ```
 
 # Original code