Update app.py

app.py

@@ -9,7 +9,7 @@ import os
 import torch.nn.functional as F
 
 num_return_sequences = 1
-max_length = 100
+max_length = 200
 
 
 @dataclass
@@ -72,7 +72,7 @@ def chat_fn(message, history):
             topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)
 
             # sampling a token from topk
-            ix = torch.multinomial(input=topk_probs, num_samples=1)  # (B, 1) (5, 1)
+            ix = torch.multinomial(input=topk_probs, num_samples=1)  # (B, 1)
 
             # gather corresponding indices
             xcol = torch.gather(input=topk_indices, dim=-1, index=ix)
@@ -83,8 +83,9 @@ def chat_fn(message, history):
         tokens = x[i, :max_length].tolist()
         decoded = tokenizer.decode(tokens)
 
-        yield decoded
+        yield decoded + "\n"
 
 
-gr.ChatInterface(chat_fn, examples=examples).launch()
-# interface.launch()
+gr.ChatInterface(chat_fn, examples=examples,
+                 title="GPT2 trained from scratch on Shakespeare dataset").launch()
+