Added output streaming support for the gradio app.

app.py

@@ -17,67 +17,78 @@ UNTRAINED_MODEL.eval()
 
 # Load fine-tuned model 
 TRAINED_MODEL = GPT(GPTConfig)
-checkpoint = torch.load("model_19072.pt", weights_only=False, map_location=torch.device('cpu'))
+checkpoint = torch.load("log/model_19072.pt", weights_only=False)
 TRAINED_MODEL.load_state_dict(checkpoint["model"])
 TRAINED_MODEL.to(device)
 TRAINED_MODEL.eval()
 
 
-def generate_text(input, model, num_sequences, max_length):
+def generate_text(input, max_length=30, top_k=50):
     tokens = TOKENIZER.encode(input)
-    tokens = torch.tensor(tokens, dtype=torch.long) 
-    tokens = tokens.unsqueeze(0).repeat(num_sequences, 1) 
-    x = tokens.to(device)
+    x_untrained = torch.tensor([tokens], dtype=torch.long).to(device)
+    x_trained   = torch.tensor([tokens], dtype=torch.long).to(device)
+
+    # Iterate until one of the sequences reaches max_length
+    while (x_untrained.size(1) < max_length) or (x_trained.size(1) < max_length):
+
+        # --- Untrained Model Forward Pass ---
+        if x_untrained.size(1) < max_length:
+            with torch.no_grad():
+                logits_u, _ = UNTRAINED_MODEL(x_untrained)
+                logits_u = logits_u[:, -1, :]
+                probs_u = F.softmax(logits_u, dim=-1)
+                topk_probs_u, topk_indices_u = torch.topk(probs_u, top_k, dim=-1)
+                ix_u = torch.multinomial(topk_probs_u, 1)
+                next_token_u = torch.gather(topk_indices_u, -1, ix_u)
+            x_untrained = torch.cat((x_untrained, next_token_u), dim=1)
+
+        # --- Trained Model Forward Pass ---
+        if x_trained.size(1) < max_length:
+            with torch.no_grad():
+                logits_t, _ = TRAINED_MODEL(x_trained)
+                logits_t = logits_t[:, -1, :]
+                probs_t = F.softmax(logits_t, dim=-1)
+                topk_probs_t, topk_indices_t = torch.topk(probs_t, top_k, dim=-1)
+                ix_t = torch.multinomial(topk_probs_t, 1)
+                next_token_t = torch.gather(topk_indices_t, -1, ix_t)
+            x_trained = torch.cat((x_trained, next_token_t), dim=1)
+
+        # --- Decode the partial text for each model ---
+        untrained_text = TOKENIZER.decode(x_untrained[0].tolist())
+        trained_text   = TOKENIZER.decode(x_trained[0].tolist())
+
+        yield (untrained_text, trained_text)
+
+
+def streaming_fn(prompt, max_length=30, top_k=50):
+    for untrained_text, trained_text in generate_text(prompt, max_length, top_k):
+        output = (
+            f"------------ (Untrained Model) ------------\n\n {untrained_text}\n\n\n"
+            f"------------ (Trained Model)------------\n\n {trained_text}"
+            )
+        yield output
+    
 
-    sentences = []
-    while x.size(1) < max_length:
-        with torch.no_grad():
-            logits, loss = model(x)
-            logits = logits[:, -1, :] 
-            probs = F.softmax(logits, dim=-1)
-            topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)
-            ix = torch.multinomial(topk_probs, 1) 
-            xcol = torch.gather(topk_indices, -1, ix) 
-            x = torch.cat((x, xcol), dim=1)
-
-    for i in range(num_sequences):
-        tokens = x[i, :max_length].tolist()
-        decoded = TOKENIZER.decode(tokens)
-        sentences.append(decoded) 
-
-    return sentences
-
-
-def gradio_fn(prompt, num_sequences=1, max_length=30):
-    """Generate text using both models."""
-    untrained_texts = generate_text(prompt, UNTRAINED_MODEL, num_sequences, max_length)
-    untrained_output = "\n\n".join(f"> {s}" for s in untrained_texts)
-
-    trained_texts = generate_text(prompt, TRAINED_MODEL, num_sequences, max_length)
-    trained_output = "\n\n".join(f"> {s}" for s in trained_texts)
-
-    return untrained_output, trained_output
-
-# Gradio interface
 def main():
     interface = gr.Interface(
-        fn=gradio_fn,
+        fn=streaming_fn,
         inputs=[
             gr.Textbox(label="Enter your prompt here:"),
-            gr.Slider(minimum=1, maximum=10, step=1, label="Number of Generations"),
-            gr.Slider(minimum=10, maximum=100, step=10, label="Max Length"),
-        ],
-        outputs=[
-            gr.Textbox(label="Generated Text (Untrained Model)"),
-            gr.Textbox(label="Generated Text (Trained Model)"),
+            gr.Slider(minimum=10, maximum=150, step=10, label="Max Length"),
+            gr.Slider(minimum=1, maximum=50, step=10, label="Top-K Samples")
         ],
-        title="GPT-2  Text Generator",
-        description="Generate text an untrained and a trained GPT-2 model."
+        outputs=gr.Textbox(label="Model Outputs"),
+        title="GPT-2 Streaming Text Generator",
+        description= (
+            "Generate text using an untrained and a trained GPT-2 model."
+            "Use prompts that are short, simple and easy to generate coherent looking text."
+            "For eg: \n"
+            "- \"Hello, my name is\" \n"
+            "- \"This is a summary of\" \n"
+            "- \"In this article\" \n"
+        )
     )
-
     interface.launch(share=True)
 
 if __name__ == "__main__":
-    main()
-
-
+    main()