Add streaming output logic

app.py

@@ -12,7 +12,8 @@ def generate_text(context, num_of_tokens, temperature=1.0):
     else:
         idx = torch.tensor(encode(context), dtype=torch.long).unsqueeze(0)
     
-    return decode(model.generate(idx, max_new_tokens=num_of_tokens,temperature=temperature)[0].tolist())
+    yield model.generate(idx, max_new_tokens=num_of_tokens,temperature=temperature)[0].tolist() 
+    
 
 
 with gr.Blocks() as demo:

model.py

@@ -206,11 +206,11 @@ class BigramLM(nn.Module):
       # sample from the distribution (pick the best)
       idx_next = torch.multinomial(probs, num_samples=1)
       # GPT like output
-      #print(decode(idx_next[0].tolist()), end='')
+      yield decode(idx_next[0].tolist())
       # append sampled index to running sequence
       idx = torch.cat((idx, idx_next), dim=1)
       
-    return idx
+    yield decode(idx_next[0].tolist())
 
 def train():