import gradio as gr
from transformers import pipeline

TITLE = "Qwen News Assistant"
DESCRIPTION = (
    "Using Qwen2.5-0.5B-Instruct to generate news drafts. "
    "This model is optimized for CPU efficiency and instruction following."
)

# 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."

    # 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(
        formatted_prompt,
        do_sample=True,
        temperature=temperature,
        top_p=float(top_p),
        max_new_tokens=int(max_length), # max_new_tokens is safer for instruct models
        pad_token_id=generator.tokenizer.eos_token_id,
    )

    # 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 = [
    ["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:
    gr.Markdown(f"# {TITLE}")
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt / Instruction",
                placeholder="e.g., 'Write a news report about...'",
                lines=4,
            )
            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 News Draft", lines=16)

    generate_button.click(
        fn=generate_text,
        inputs=[prompt, temperature, top_p, max_length],
        outputs=output,
    )

    gr.Examples(
        examples=examples,
        inputs=[prompt, temperature, top_p, max_length],
        outputs=output,
        fn=generate_text,
        cache_examples=False,
    )

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