Updated README

Author: rish-16

README.md

@@ -2,13 +2,14 @@
 
 <p align="center">Easy-to-use Wrapper for GPT-2 117M and 345M Transformer Models</p>
 
-<!-- <center>
+<div align="center" style="display: flex;justify-content:space-evenly;width: 250px;margin: 0 auto;">
 
-[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+<p align="center"><a style="margin: 0 10px" href="https://opensource.org/licenses/MIT"><img src="https://img.shields.io/badge/License-MIT-yellow.svg"></a></p>
+<p align="center"><a style="margin: 0 10px" href="https://colab.research.google.com/github/googlecolab/colabtools/blob/master/notebooks/colab-github-demo.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg"></a></p>
 
-</center> -->
+</div>
 
-<img src="https://github.com/rish-16/gpt2client/raw/master/assets/demo.png" style="width: 80%;" />
+<div><img src="https://github.com/rish-16/gpt2client/raw/master/assets/demo.png" /></div>
 
 <p align="center"><strong>Made by Rishabh Anand • <a href="https://rish-16.github.io">https://rish-16.github.io</a></strong></p>
 

gpt2_client/gpt2_client.py

@@ -178,6 +178,7 @@ def generate(self, interactive=False, n_samples=1, words=None, display=True, ret
 					context_tokens = enc.encode(prompt)
 					text_array = []
 					text = ''
+					generated = 0
 					for _ in range(n_samples // batch_size):
 						out = sess.run(output, feed_dict={
 							context: [context_tokens for _ in range(batch_size)]