from transformers import pipeline, set_seed
import gradio as gr
import re

generator = pipeline(
    "text-generation",
    model="Keyurjotaniya007/openai-gpt2-wikitext-textgen"
)

set_seed(42)

def clean_output(text):
    text = re.sub(r"=+\s*[^=]+?\s*=+", "", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text

def generate_text(prompt):
    if not prompt.strip():
        return "Please enter a prompt to generate text."
    response = generator(
        prompt,
        max_length=150,
        num_return_sequences=1,
        temperature=0.8,
        top_p=0.95,
        do_sample=True,
        truncation=True
    )
    return response[0]["generated_text"]

interface = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(lines=4, placeholder="Type your prompt here..."),
    outputs=gr.Textbox(label="Generated Text"),
    title="🦄 GPT-2 WikiText Text Generator",
    description=(
        "Generate coherent, fluent English text using **GPT-2 fine-tuned on the WikiText dataset**. "
        "Model: `Keyurjotaniya007/openai-gpt2-wikitext-textgen`"
    ),
    examples=[
        ["Artificial intelligence is transforming the world by"],
        ["Once upon a time, in a small village,"],
    ],
)

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