Suzana commited on
Commit
4473874
·
verified ·
1 Parent(s): 6cfea42

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +91 -0
app.py ADDED
@@ -0,0 +1,91 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import pipeline
3
+ import os
4
+
5
+ # Cache pipelines so we don’t reload every call
6
+ pipeline_cache = {}
7
+
8
+ def get_generator(model_name):
9
+ if model_name not in pipeline_cache:
10
+ pipeline_cache[model_name] = pipeline(
11
+ "text-generation",
12
+ model=model_name,
13
+ device=0 if os.environ.get("CUDA_VISIBLE_DEVICES") else -1
14
+ )
15
+ return pipeline_cache[model_name]
16
+
17
+ def build_prompt(system, examples, query, style):
18
+ # Prefix with system role
19
+ prompt = f"SYSTEM: {system.strip()}\n\n"
20
+ if style == "Zero-Shot":
21
+ prompt += f"USER: {query.strip()}\nASSISTANT:"
22
+ else: # Few-Shot
23
+ # Insert up to three examples
24
+ for ex in examples:
25
+ ex = ex.strip()
26
+ if ex:
27
+ prompt += f"EXAMPLE: {ex}\n"
28
+ prompt += f"USER: {query.strip()}\nASSISTANT:"
29
+ return prompt
30
+
31
+ def generate(system, ex1, ex2, ex3, query, model_name, style):
32
+ # Build prompt text
33
+ examples = [ex1, ex2, ex3]
34
+ prompt = build_prompt(system, examples, query, style)
35
+ gen = get_generator(model_name)
36
+ # Generate up to 100 tokens beyond prompt
37
+ out = gen(prompt, max_length=len(prompt.split())+100, do_sample=True)[0]["generated_text"]
38
+ # Strip prompt prefix
39
+ return out[len(prompt):].strip()
40
+
41
+ # Available small/medium models on HF Hub
42
+ AVAILABLE_MODELS = [
43
+ "gpt2",
44
+ "EleutherAI/gpt-neo-125M",
45
+ "distilgpt2"
46
+ ]
47
+
48
+ with gr.Blocks(theme=gr.themes.Default()) as demo:
49
+ gr.Markdown("## 🎮 Prompt Engineering Playground")
50
+ gr.Markdown(
51
+ "Experiment with **Zero-Shot** vs **Few-Shot** prompting across different open models."
52
+ )
53
+
54
+ with gr.Row():
55
+ with gr.Column(scale=2):
56
+ system_box = gr.Textbox(
57
+ label="🖥️ System Prompt",
58
+ placeholder="You are a helpful assistant."
59
+ )
60
+ ex1 = gr.Textbox(label="📖 Example 1 (few-shot only)", placeholder="…")
61
+ ex2 = gr.Textbox(label="📖 Example 2 (few-shot only)")
62
+ ex3 = gr.Textbox(label="📖 Example 3 (few-shot only)")
63
+ query_box = gr.Textbox(label="🗣️ User Prompt", placeholder="Type your query here…")
64
+
65
+ with gr.Row():
66
+ model_choice = gr.Dropdown(
67
+ choices=AVAILABLE_MODELS,
68
+ value=AVAILABLE_MODELS[0],
69
+ label="🤖 Model"
70
+ )
71
+ style_choice = gr.Radio(
72
+ choices=["Zero-Shot", "Few-Shot"],
73
+ value="Zero-Shot",
74
+ label="🎨 Prompt Style"
75
+ )
76
+
77
+ run_btn = gr.Button("🚀 Generate")
78
+
79
+ with gr.Column(scale=3):
80
+ output_box = gr.Textbox(
81
+ label="💬 Model Output",
82
+ lines=15
83
+ )
84
+
85
+ run_btn.click(
86
+ fn=generate,
87
+ inputs=[system_box, ex1, ex2, ex3, query_box, model_choice, style_choice],
88
+ outputs=output_box
89
+ )
90
+
91
+ demo.launch()