import gradio as gr from transformers import pipeline import os # Cache pipelines so we donโ€™t reload every call pipeline_cache = {} def get_generator(model_name): if model_name not in pipeline_cache: pipeline_cache[model_name] = pipeline( "text-generation", model=model_name, device=0 if os.environ.get("CUDA_VISIBLE_DEVICES") else -1 ) return pipeline_cache[model_name] def build_prompt(system, examples, query, style): # Prefix with system role prompt = f"SYSTEM: {system.strip()}\n\n" if style == "Zero-Shot": prompt += f"USER: {query.strip()}\nASSISTANT:" else: # Few-Shot # Insert up to three examples for ex in examples: ex = ex.strip() if ex: prompt += f"EXAMPLE: {ex}\n" prompt += f"USER: {query.strip()}\nASSISTANT:" return prompt def generate(system, ex1, ex2, ex3, query, model_name, style): # Build prompt text examples = [ex1, ex2, ex3] prompt = build_prompt(system, examples, query, style) gen = get_generator(model_name) # Generate up to 100 tokens beyond prompt out = gen(prompt, max_length=len(prompt.split())+100, do_sample=True)[0]["generated_text"] # Strip prompt prefix return out[len(prompt):].strip() # Available small/medium models on HF Hub AVAILABLE_MODELS = [ "gpt2", "EleutherAI/gpt-neo-125M", "distilgpt2" ] with gr.Blocks(theme=gr.themes.Default()) as demo: gr.Markdown("## ๐ŸŽฎ Prompt Engineering Playground") gr.Markdown( "Experiment with **Zero-Shot** vs **Few-Shot** prompting across different open models." ) with gr.Row(): with gr.Column(scale=2): system_box = gr.Textbox( label="๐Ÿ–ฅ๏ธ System Prompt", placeholder="You are a helpful assistant." ) ex1 = gr.Textbox(label="๐Ÿ“– Example 1 (few-shot only)", placeholder="โ€ฆ") ex2 = gr.Textbox(label="๐Ÿ“– Example 2 (few-shot only)") ex3 = gr.Textbox(label="๐Ÿ“– Example 3 (few-shot only)") query_box = gr.Textbox(label="๐Ÿ—ฃ๏ธ User Prompt", placeholder="Type your query hereโ€ฆ") with gr.Row(): model_choice = gr.Dropdown( choices=AVAILABLE_MODELS, value=AVAILABLE_MODELS[0], label="๐Ÿค– Model" ) style_choice = gr.Radio( choices=["Zero-Shot", "Few-Shot"], value="Zero-Shot", label="๐ŸŽจ Prompt Style" ) run_btn = gr.Button("๐Ÿš€ Generate") with gr.Column(scale=3): output_box = gr.Textbox( label="๐Ÿ’ฌ Model Output", lines=15 ) run_btn.click( fn=generate, inputs=[system_box, ex1, ex2, ex3, query_box, model_choice, style_choice], outputs=output_box ) demo.launch()