Revert streaming logic

app.py

@@ -6,13 +6,13 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 model = torch.load('saved_model.pth', map_location= torch.device(device), weights_only=False)
 
-# def generate_text(context, num_of_tokens, temperature=1.0):
-#     if context == None or context == '':
-#         idx = torch.zeros((1,1), dtype=torch.long)
-#     else:
-#         idx = torch.tensor(encode(context), dtype=torch.long).unsqueeze(0)
+def generate_text(context, num_of_tokens, temperature=1.0):
+    if context == None or context == '':
+        idx = torch.zeros((1,1), dtype=torch.long)
+    else:
+        idx = torch.tensor(encode(context), dtype=torch.long).unsqueeze(0)
     
-#     yield model.generate(idx, max_new_tokens=num_of_tokens,temperature=temperature) 
+    yield model.generate(idx, max_new_tokens=num_of_tokens,temperature=temperature) 
     
 
 
@@ -29,10 +29,6 @@ with gr.Blocks() as demo:
         context,
         num_of_tokens,tmp
     ]
-    if context == None or context == '':
-        idx = torch.zeros((1,1), dtype=torch.long)
-    else:
-        idx = torch.tensor(encode(context), dtype=torch.long).unsqueeze(0)
     generate_btn = gr.Button(value="Generate")
     outputs = [gr.Textbox(label= "Generated text: ")]
     generate_btn.click(fn = model.generate, inputs= inputs, outputs= outputs)

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
-      yield decode(idx_next[0].tolist())
+      # yield decode(idx_next[0].tolist())
       # append sampled index to running sequence
       idx = torch.cat((idx, idx_next), dim=1)
       
-    yield decode(idx_next[0].tolist())
+    return idx
 
 def train():