import gradio as gr
from transformers import pipeline

# Initialize the text generation pipeline
generator = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1", trust_remote_code=True)

def generate_text(prompt):
    try:
        # Generate text with a maximum length of 50 tokens
        output = generator(prompt, max_length=50)
        generated_text = output[0]['generated_text']
        return generated_text
    except Exception as e:
        return f"Error: {str(e)}"

# Define the Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your text here..."),
    outputs="text",
    title="Text Generation with Hugging Face Transformers",
    description="Enter a prompt and generate text using the DeepSeek-R1 model."
)

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