from transformers import pipeline
import torch
import gradio as gr

pipeline = pipeline("text-generation", model = "distilbert/distilgpt2")
def generate_text(text:str):
    """Complete the text
    Input:
            Text : Given some Initial word.
    Output:
          Complete the text till the max length.
    """
    return pipeline(text, max_length=50, truncation=True)

demo = gr.Interface(
    generate_text,
    inputs=["text"],
    outputs=["json"]
)

demo.launch(mcp_server = True)