Update app.py

app.py

@@ -13,47 +13,35 @@ model = GPT2LMHeadModel.from_pretrained(model_path).to(device)
 # Set model to evaluation mode
 model.eval()
 
-# Function to generate text in a stream-based manner
+# Function to generate text based on input prompt
 def generate_text(prompt):
     # Tokenize and encode the input prompt
     input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
-    max_length = 50  # Maximum length of generated text
-
-    # Generate continuation with streaming tokens
-    output_ids = input_ids
-    for _ in range(max_length):
-        # Generate the next token only (without generating full sequence)
-        with torch.no_grad():
-            next_token_ids = model.generate(
-                output_ids, 
-                max_length=output_ids.shape[-1] + 1,  # Increase length by 1 each step
-                num_return_sequences=1, 
-                pad_token_id=tokenizer.eos_token_id, 
-                do_sample=True, 
-                top_k=50, 
-                top_p=0.95,
-                use_cache=True
-            )
-
-        # Get the newly generated token (last one in the sequence)
-        next_token = next_token_ids[:, -1:]
-
-        # Append new token to the output sequence
-        output_ids = torch.cat((output_ids, next_token), dim=-1)
-
-        # Decode and yield the output text incrementally
-        decoded_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-        yield decoded_text
-
-# Create a Gradio interface with streaming enabled
+
+    # Generate continuation
+    with torch.no_grad():
+        generated_ids = model.generate(
+            input_ids, 
+            max_length=50,                # Maximum length of generated text
+            num_return_sequences=1,        # Generate 1 sequence
+            pad_token_id=tokenizer.eos_token_id,  # Use EOS token for padding
+            do_sample=True,                # Enable sampling
+            top_k=50,                      # Top-k sampling
+            top_p=0.95                     # Nucleus sampling
+        )
+
+    # Decode the generated text
+    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    return generated_text
+
+# Create a Gradio interface
 interface = gr.Interface(
-    fn=generate_text,                          # Function to call when interacting with the UI
-    inputs="text",                             # Input type: Single-line text
-    outputs=gr.Textbox(),                      # Stream output using a Textbox for real-time updates
-    title="Quble Text Generation",             # Title of the UI
-    description="Enter a prompt to generate text using Quble with live streaming.",  # Simple description
-    live=True                                  # Enable live streaming of the output
+    fn=generate_text,                      # Function to call when interacting with the UI
+    inputs="text",                         # Input type: Single-line text
+    outputs="text",                        # Output type: Text (the generated output)
+    title="Quble Text Generation",         # Title of the UI
+    description="Enter a prompt to generate text using Quble."  # Simple description
 )
 
 # Launch the Gradio app
-interface.launch()
+interface.launch()