Update app.py

app.py

@@ -25,7 +25,7 @@ def generate(
             temperature = 0.9,
             verbose=False,
             return_prime=False,
-            ):
+            progress=gr.Progress()):
 
     out = torch.LongTensor([start_tokens])
 
@@ -34,29 +34,23 @@ def generate(
     if verbose:
       print("Generating sequence of max length:", seq_len)
 
-    max_len = seq_len
-    cur_len = 0
-        
-    bar = tqdm.tqdm(desc="generating", total=max_len - cur_len, disable=in_space)
-    with bar:
-        while cur_len < max_len:
-    
-            x = out[:, -max_seq_len:]
-    
-            torch_in = x.tolist()[0]
-    
-            logits = torch.FloatTensor(session.run(None, {'input': [torch_in]})[0])[:, -1]
-            
-            filtered_logits = logits
-    
-            probs = F.softmax(filtered_logits / temperature, dim=-1)
-    
-            sample = torch.multinomial(probs, 1)
+    progress(0, desc="Starting...")
+                
+    for i in progress.tqdm(range(seq_len)):
     
-            out = torch.cat((out, sample), dim=-1)
-            
-            cur_len += 1
-            bar.update(1)
+        x = out[:, -max_seq_len:]
+
+        torch_in = x.tolist()[0]
+
+        logits = torch.FloatTensor(session.run(None, {'input': [torch_in]})[0])[:, -1]
+        
+        filtered_logits = logits
+
+        probs = F.softmax(filtered_logits / temperature, dim=-1)
+
+        sample = torch.multinomial(probs, 1)
+
+        out = torch.cat((out, sample), dim=-1)
             
     if return_prime:
       return out[:, :]