kgourgou commited on
Commit
0aca33c
·
verified ·
1 Parent(s): dd1fb47

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +44 -0
app.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+
4
+ # Load your model (using GPT-2 as an example)
5
+ model_name = "gpt2"
6
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
7
+ model = AutoModelForCausalLM.from_pretrained(model_name)
8
+
9
+ def generate_completions(prompt):
10
+ # Define decoding strategies with corresponding parameters
11
+ strategies = {
12
+ "Greedy": {"do_sample": False},
13
+ "Beam Search": {"num_beams": 5, "early_stopping": True},
14
+ "Top-k Sampling": {"do_sample": True, "top_k": 50},
15
+ "Top-p Sampling": {"do_sample": True, "top_p": 0.95}
16
+ }
17
+
18
+ results = {}
19
+ input_ids = tokenizer.encode(prompt, return_tensors="pt")
20
+
21
+ for strategy, params in strategies.items():
22
+ # Generate output using the specific strategy
23
+ output_ids = model.generate(input_ids, max_length=50, **params)
24
+ output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
25
+ results[strategy] = output_text
26
+
27
+ return results["Greedy"], results["Beam Search"], results["Top-k Sampling"], results["Top-p Sampling"]
28
+
29
+ # Define the Gradio interface
30
+ interface = gr.Interface(
31
+ fn=generate_completions,
32
+ inputs=gr.inputs.Textbox(lines=3, placeholder="Enter your prompt here...", label="Prompt"),
33
+ outputs=[
34
+ gr.outputs.Textbox(label="Greedy"),
35
+ gr.outputs.Textbox(label="Beam Search"),
36
+ gr.outputs.Textbox(label="Top-k Sampling"),
37
+ gr.outputs.Textbox(label="Top-p Sampling"),
38
+ ],
39
+ title="LLM Decoding Strategies Comparison",
40
+ description="Enter a prompt to see how different decoding strategies affect the output of a language model."
41
+ )
42
+
43
+ if __name__ == "__main__":
44
+ interface.launch()