Spaces:
Sleeping
Sleeping
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()
|