Update app.py

app.py

@@ -1,44 +1,91 @@
 import gradio as gr
-from transformers import pipeline
+from functools import lru_cache
 
-# 使用一个轻量文本生成模型（CPU 可跑）
-generator = pipeline(
-    "text-generation",
-    model="distilgpt2"
-)
+@lru_cache(maxsize=1)
+def get_rewriter():
+    """
+    Use a lightweight instruction-following model that can run on CPU.
+    flan-t5-small is generally more reliable for "rewrite" than GPT2-style models.
+    """
+    from transformers import pipeline
+    return pipeline(
+        task="text2text-generation",
+        model="google/flan-t5-small",
+        device=-1,  # CPU
+    )
+
+def build_prompt(text: str, style: str) -> str:
+    text = (text or "").strip()
+    if not text:
+        return ""
 
-def rewrite(text, style):
     if style == "More formal":
-        prompt = f"Rewrite this sentence in a more formal tone:\n{text}"
-    elif style == "More friendly":
-        prompt = f"Rewrite this sentence in a friendly tone:\n{text}"
-    else:
-        prompt = f"Rewrite this sentence to be shorter:\n{text}"
+        return (
+            "Rewrite the text in a more formal tone. "
+            "Keep the original meaning. Output only the rewritten text.\n\n"
+            f"Text: {text}"
+        )
+    if style == "More friendly":
+        return (
+            "Rewrite the text in a friendly, warm tone. "
+            "Keep the original meaning. Output only the rewritten text.\n\n"
+            f"Text: {text}"
+        )
+    return (
+        "Rewrite the text to be shorter and clearer. "
+        "Keep the original meaning. Output only the rewritten text.\n\n"
+        f"Text: {text}"
+    )
+
+def rewrite(text: str, style: str) -> str:
+    text = (text or "").strip()
+    if not text:
+        return "Please enter some text."
 
-    result = generator(prompt, max_length=100, num_return_sequences=1)
-    return result[0]["generated_text"]
+    prompt = build_prompt(text, style)
+
+    try:
+        rewriter = get_rewriter()
+        out = rewriter(
+            prompt,
+            max_new_tokens=128,
+            do_sample=False,
+        )
+        result = (out[0].get("generated_text") or "").strip()
+        return result if result else "No output. Try a shorter input."
+    except Exception as e:
+        return (
+            "Error: failed to run the model.\n"
+            "If this is the first run, the Space may still be downloading the model.\n\n"
+            f"Details: {type(e).__name__}: {e}"
+        )
 
 with gr.Blocks(title="AI Text Rewriter") as demo:
     gr.Markdown(
-        "### ✨ AI Text Rewriter\n"
-        "Paste a sentence and let AI rewrite it for you."
+        "AI Text Rewriter\n"
+        "Paste a sentence or short paragraph, choose a style, then rewrite with AI."
     )
 
+    with gr.Row():
+        style = gr.Radio(
+            ["More formal", "More friendly", "Shorter"],
+            value="More friendly",
+            label="Rewrite style",
+        )
+
     text_input = gr.Textbox(
         label="Your text",
-        lines=3,
-        placeholder="Type something here..."
+        placeholder="Type or paste text here...",
+        lines=5,
     )
 
-    style = gr.Radio(
-        ["More formal", "More friendly", "Shorter"],
-        label="Rewrite style",
-        value="More friendly"
-    )
+    with gr.Row():
+        btn = gr.Button("Rewrite with AI")
+        clear = gr.Button("Clear")
 
-    output = gr.Textbox(label="AI Result", lines=4)
+    output = gr.Textbox(label="Result", lines=6)
 
-    btn = gr.Button("Rewrite with AI")
     btn.click(fn=rewrite, inputs=[text_input, style], outputs=output)
+    clear.click(fn=lambda: ("", ""), inputs=None, outputs=[text_input, output])
 
 demo.launch()