Update app.py

app.py

@@ -1,41 +1,53 @@
 import gradio as gr
 from transformers import pipeline
 
-TITLE = "News Draft Generator"
+TITLE = "Qwen News Assistant"
 DESCRIPTION = (
-    "Experiment with news-style writing using distilgpt2. "
-    "Try headlines, article openings, research summaries, and briefings."
+    "Using Qwen2.5-0.5B-Instruct to generate news drafts. "
+    "This model is optimized for CPU efficiency and instruction following."
 )
 
-generator = pipeline("text-generation", model="distilgpt2")
-
+# Load the lightweight Qwen 0.5B model
+generator = pipeline("text-generation", model="Qwen/Qwen2.5-0.5B-Instruct")
 
 def generate_text(prompt: str, temperature: float, top_p: float, max_length: int) -> str:
     if not prompt or not prompt.strip():
         return "Please enter a prompt."
 
-    # Clamp temperature to avoid division-by-zero-style issues in downstream sampling
+    # Instruct models perform best when prompts are formatted correctly
+    # For Qwen, we wrap the prompt in a simple instruction format
+    messages = [
+        {"role": "system", "content": "You are a helpful news writing assistant."},
+        {"role": "user", "content": prompt},
+    ]
+    
+    # Apply the chat template
+    formatted_prompt = generator.tokenizer.apply_chat_template(
+        messages, 
+        tokenize=False, 
+        add_generation_prompt=True
+    )
+
     temperature = max(0.01, float(temperature))
 
     outputs = generator(
-        prompt,
+        formatted_prompt,
         do_sample=True,
         temperature=temperature,
         top_p=float(top_p),
-        max_length=int(max_length),
-        num_return_sequences=1,
+        max_new_tokens=int(max_length), # max_new_tokens is safer for instruct models
         pad_token_id=generator.tokenizer.eos_token_id,
     )
 
-    return outputs[0]["generated_text"]
-
+    # Extract only the newly generated text
+    generated_text = outputs[0]["generated_text"]
+    return generated_text.split("<|im_start|>assistant\n")[-1].strip()
 
+# Updated examples to be more "Instruction" focused
 examples = [
-    ["Breaking news: scientists discovered that", 0.5, 0.9, 100],
-    ["A new report released today found that students", 0.5, 0.9, 100],
-    ["The conference on artificial intelligence announced", 0.5, 0.9, 100],
-    ["Researchers at the university published a study showing", 0.5, 0.9, 100],
-    ["This week in technology: the biggest story is", 0.5, 0.9, 100],
+    ["Write a breaking news headline about a discovery on Mars.", 0.7, 0.9, 100],
+    ["Write the opening paragraph for a story about local students winning a robotics competition.", 0.5, 0.9, 150],
+    ["Summarize the importance of artificial intelligence in 2024.", 0.3, 0.9, 200],
 ]
 
 with gr.Blocks(title=TITLE) as demo:
@@ -45,45 +57,17 @@ with gr.Blocks(title=TITLE) as demo:
     with gr.Row():
         with gr.Column():
             prompt = gr.Textbox(
-                label="Prompt",
-                placeholder="Enter a headline, opening sentence, or briefing prompt...",
+                label="Prompt / Instruction",
+                placeholder="e.g., 'Write a news report about...'",
                 lines=4,
             )
-
-            temperature = gr.Slider(
-                minimum=0.1,
-                maximum=2.0,
-                value=0.5,
-                step=0.1,
-                label="Temperature",
-                info="Controls how creative/wild the writing is",
-            )
-
-            top_p = gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=0.9,
-                step=0.05,
-                label="Top-p",
-                info="Controls word diversity",
-            )
-
-            max_length = gr.Slider(
-                minimum=20,
-                maximum=200,
-                value=100,
-                step=1,
-                label="Max Length",
-                info="Controls how much text it generates",
-            )
-
-            generate_button = gr.Button("Generate")
+            temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
+            top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p")
+            max_length = gr.Slider(minimum=20, maximum=300, value=150, step=1, label="Max New Tokens")
+            generate_button = gr.Button("Generate", variant="primary")
 
         with gr.Column():
-            output = gr.Textbox(
-                label="Generated Text",
-                lines=16,
-            )
+            output = gr.Textbox(label="Generated News Draft", lines=16)
 
     generate_button.click(
         fn=generate_text,