Spaces:
Sleeping
Sleeping
File size: 4,294 Bytes
6cdbea4 abd0d11 6cdbea4 |
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 93 94 |
import time
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
# Load model and tokenizer once at startup
model_name = "Qwen/Qwen2.5-0.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
print(f"Model loaded on: {model.device}")
# Define the generation function
def generate_text(prompt, max_new_tokens, num_beams):
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
input_length = inputs["input_ids"].shape[-1]
# Greedy search
start_time = time.time()
outputs_greedy = model.generate(
**inputs,
max_new_tokens=int(max_new_tokens),
num_beams=1,
do_sample=False,
)
greedy_time = time.time() - start_time
# Remove the prompt tokens from the output
generated_tokens_greedy = outputs_greedy[0][input_length:]
generated_text_greedy = tokenizer.decode(generated_tokens_greedy, skip_special_tokens=True)
# Beam search
start_time = time.time()
outputs_beam = model.generate(
**inputs,
num_beams=int(num_beams),
num_return_sequences=1,
max_new_tokens=int(max_new_tokens),
do_sample=False,
)
beam_time = time.time() - start_time
# Remove the prompt tokens as above
generated_tokens_beam = outputs_beam[0][input_length:]
generated_text_beam = tokenizer.decode(generated_tokens_beam, skip_special_tokens=True)
# Prepare outputs for better display formatting
greedy_details = (
f"**Strategy:** Picks the most probable token at each step (deterministic).\n\n"
f"**Time:** {greedy_time:.2f} seconds"
)
beam_details = (
f"**Strategy:** Explores {num_beams} beams concurrently and returns the top candidate.\n\n"
f"**Time:** {beam_time:.2f} seconds"
)
return greedy_details, generated_text_greedy, beam_details, generated_text_beam
with gr.Blocks() as demo:
# Informational header to help users understand the demo
gr.Markdown(
"# Beam Search Demo\n\n"
"This demo shows how two different text generation strategies work using the Qwen2.5-0.5B model. "
"The left side uses **greedy search**, which picks the most probable token at every generation step (deterministic), "
"while the right side uses **beam search**, which explores multiple beams concurrently to choose the most likely "
"sequence of tokens.\n\n"
"**Important:** This model works best with prompts that need completion rather than question-answering. For example, "
"instead of 'What is the capital of France?', use prompts like 'The capital of France is' or 'Here is a story about:'\n\n"
"Use the controls below to enter your prompt, adjust the maximum number of newly generated tokens, and set the "
"number of beams for beam search. The results for both strategies are displayed side by side for easy comparison.\n\n"
"Repo: [Beam Search Demo](https://github.com/cavit99/beam-search-demo)"
)
# Input components in a single column at the top
with gr.Column():
gr.Markdown("## Input")
prompt_input = gr.Textbox(label="Prompt", value="Here is a funny love letter for you:")
max_tokens_input = gr.Slider(minimum=1, maximum=100, step=1, label="Max new tokens", value=50)
num_beams_input = gr.Slider(minimum=1, maximum=20, step=1, label="Number of beams", value=10)
generate_btn = gr.Button("Generate")
with gr.Row():
with gr.Column():
greedy_details_output = gr.Markdown(label="Greedy Search Details")
greedy_textbox_output = gr.Textbox(label="Greedy Search Generated Text", lines=10)
with gr.Column():
beam_details_output = gr.Markdown(label="Beam Search Details")
beam_textbox_output = gr.Textbox(label="Beam Search Generated Text", lines=10)
# Connect the button click event to the generation function
generate_btn.click(
generate_text,
inputs=[prompt_input, max_tokens_input, num_beams_input],
outputs=[greedy_details_output, greedy_textbox_output, beam_details_output, beam_textbox_output]
)
if __name__ == "__main__":
demo.launch() |