Update app.py

app.py

@@ -2,10 +2,10 @@ import gradio as gr
 import torch
 from mingru_lm import MinGRU_LM
 
-
+# Load the model
 model = MinGRU_LM(dim=512, num_tokens=256, num_layers=6)
 pt_model = "best_model.pt"
-checkpoint = torch.load(pt_model,map_location=torch.device('cpu'))
+checkpoint = torch.load(pt_model, map_location=torch.device('cpu'))
 model.load_state_dict(checkpoint['model_state_dict'])
 
 # Move model to GPU if available
@@ -25,36 +25,52 @@ def generate_text(start_text, max_length, temperature):
     input_tensor = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(device)  # Ensure long tensor
     
     generated_tokens = tokens.copy()
-    
-    with torch.no_grad():
-        for _ in range(max_length):
-            _, logits = model(input_tensor, labels=None)
+
+    # Use a generator to yield tokens one by one
+    for _ in range(max_length):
+        with torch.no_grad():
+            logits = model(input_tensor, labels=None)[1]  # Get logits directly
             
             last_token_logits = logits[0, -1, :] / temperature
             probs = torch.softmax(last_token_logits, dim=-1)
+            
+            # Sample the next token
             next_token = torch.multinomial(probs, num_samples=1).item()
             
-            # Only append if it's within the 256-character ASCII range
+            # Only append valid tokens
             if next_token < 256:  
                 generated_tokens.append(next_token)
                 input_tensor = torch.cat([input_tensor, torch.tensor([[next_token]], device=device)], dim=1)
+                
+                yield decode_tokens(generated_tokens)
             else:
                 continue  # Skip tokens outside ASCII range
         
-    return decode_tokens(generated_tokens)
+    yield decode_tokens(generated_tokens)
+
+def wrapper_generate_text(start_text, max_length, temperature):
+    async_gen = generate_text(start_text, max_length, temperature)
+    
+    for output in async_gen:
+        yield output
 
 # Gradio interface
-iface = gr.Interface(
-    fn=generate_text,
-    inputs=[
-        gr.Textbox(lines=3, label="Enter your prompt", value="Once upon a time"),
-        gr.Slider(minimum=10, maximum=500, value=200, step=1, label="Max Length"),
-        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
-    ],
-    outputs=gr.Textbox(lines=10, label="Generated Text"),
-    title="Text Generation with MinGRU_LM",
-    description="Enter a prompt and adjust parameters to generate text using the MinGRU_LM model."
-)
-
-if __name__ == "__main__":
-    iface.launch(show_api=False, server_name="0.0.0.0")
+with gr.Blocks() as iface:
+    gr.Markdown("### Please be patient, generating text will take some time...")
+    
+    with gr.Row():
+        textbox = gr.Textbox(lines=3, label="Enter your prompt", value="Once upon a time")
+        max_length = gr.Slider(minimum=10, maximum=500, value=200, step=1, label="Max Length")
+        temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
+    
+    output_textbox = gr.Textbox(lines=10, label="Generated Text")
+    
+    btn = gr.Button("Generate Text")
+    
+    btn.click(
+        wrapper_generate_text, 
+        inputs=[textbox, max_length, temperature],
+        outputs=output_textbox
+    )
+
+iface.launch(show_api=False, server_name="0.0.0.0")