import gradio as gr
from generate import generate_text
import torch

def generate(prompt):
    try:
        # Redirect print output to capture the generated text
        import io
        import sys
        old_stdout = sys.stdout
        new_stdout = io.StringIO()
        sys.stdout = new_stdout
        
        # Generate text with default parameters
        generate_text(
            prompt=prompt,
            max_length=100,  # default value
            num_sequences=5  # default value
        )
        
        # Get the output and restore stdout
        output = new_stdout.getvalue()
        sys.stdout = old_stdout
        
        return output
    except Exception as e:
        return f"An error occurred: {str(e)}"

# Create the Gradio interface
demo = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(label="Enter your prompt", lines=3),
    outputs=gr.Textbox(label="Generated Text", lines=10),
    title="Text Generation Demo",
    description="Enter a prompt and get generated text responses.",
)

if __name__ == "__main__":
    demo.launch()