File size: 3,044 Bytes
4473874
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
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()