joaogante HF staff commited on
Commit
9015a60
1 Parent(s): d074cd5

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +56 -46
app.py CHANGED
@@ -174,9 +174,6 @@ STYLE = """
174
  .nonselected-sequence {
175
  background-color: var(--primary-500);
176
  }
177
- .nopadding {
178
- padding-left: 0!important;
179
- }
180
  """
181
 
182
 
@@ -241,7 +238,7 @@ def generate_nodes(node, step):
241
  def generate_html(start_sentence, original_tree):
242
  html_output = f"""<div class="custom-container">
243
  <div class="tree">
244
- <ul class="nopadding"> <li class="nopadding"> <a id='root'> <span> <b>{start_sentence}</b> </span> {original_tree.table} </a>"""
245
  html_output += "<ul> "
246
  for subnode in original_tree.children.values():
247
  html_output += generate_nodes(subnode, step=1)
@@ -273,20 +270,20 @@ class BeamNode:
273
 
274
 
275
  def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, beam_indexes_source):
276
- input_length = len(tokenizer([start_sentence], return_tensors="pt"))
277
  original_tree = BeamNode(
278
  cumulative_score=0,
279
  current_token_ix=None,
280
  table=None,
281
  current_sequence=start_sentence,
282
  children={},
283
- children_score_divider=((input_length + 1) ** length_penalty),
284
  total_score=None,
285
  is_final=False,
286
  is_selected_sequence=False,
287
  )
288
  n_beams = len(scores[0])
289
  beam_trees = [original_tree] * n_beams
 
290
 
291
  for step, step_scores in enumerate(scores):
292
 
@@ -297,8 +294,11 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, be
297
  beam_indexes,
298
  current_sequence,
299
  top_tokens,
300
- ) = ([], [], [], [], [])
301
- for beam_ix in range(n_beams):
 
 
 
302
  current_beam = beam_trees[beam_ix]
303
 
304
  # skip if the beam is already final
@@ -307,16 +307,18 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, be
307
 
308
  # Get top cumulative scores for the current beam
309
  current_top_token_indexes = list(
310
- np.array(scores[step][beam_ix].argsort()[-n_beams:])[::-1]
311
  )
312
  top_token_indexes += current_top_token_indexes
 
313
  top_cumulative_scores += list(
314
- np.array(scores[step][beam_ix][current_top_token_indexes])
315
  + current_beam.cumulative_score
316
  )
317
  beam_indexes += [beam_ix] * n_beams
318
  current_sequence += [beam_trees[beam_ix].current_sequence] * n_beams
319
  top_tokens += [tokenizer.decode([el]) for el in current_top_token_indexes]
 
320
 
321
  top_df = pd.DataFrame.from_dict(
322
  {
@@ -325,6 +327,7 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, be
325
  "beam_index": beam_indexes,
326
  "current_sequence": current_sequence,
327
  "token": top_tokens,
 
328
  }
329
  )
330
  maxes = top_df.groupby(["token_index", "current_sequence"])[
@@ -333,78 +336,85 @@ def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, be
333
 
334
  top_df = top_df.loc[maxes]
335
 
336
- # Sort all top probabilities and keep top n_beams
 
337
  top_df_selected = top_df.sort_values("cumulative_score", ascending=False).iloc[
338
- :n_beams
339
  ]
340
- if any(["you enjoyed" in el for el in top_df["current_sequence"]]):
341
- print("Displaying debug info:::")
342
- print(top_df_selected)
 
 
 
 
 
 
 
 
 
 
343
 
344
  # Write the scores table - one per beam source
345
- for beam_ix in reversed(list(range(n_beams))):
 
346
  current_beam = beam_trees[beam_ix]
347
  if current_beam.table is None:
348
- selected_tokens = top_df_selected.loc[
349
- top_df_selected["current_sequence"] == current_beam.current_sequence
350
  ]
351
  markdown_table = generate_markdown_table(
352
- step_scores[beam_ix, :],
353
  current_beam.cumulative_score,
354
  current_beam.children_score_divider,
355
  chosen_tokens=list(selected_tokens["token"].values),
356
  )
357
  beam_trees[beam_ix].table = markdown_table
 
 
358
 
359
  # Add new children to each beam
360
  cumulative_scores = [beam.cumulative_score for beam in beam_trees]
361
- for _, row in top_df_selected.iterrows():
362
  # Update the source tree
363
  source_beam_ix = int(row["beam_index"])
364
  current_token_choice_ix = row["token_index"]
365
  current_token_choice = tokenizer.decode([current_token_choice_ix])
 
366
 
367
- cumulative_score = (
368
- cumulative_scores[source_beam_ix]
369
- + scores[step][source_beam_ix][current_token_choice_ix].numpy()
370
- )
371
  current_sequence = (
372
  beam_trees[source_beam_ix].current_sequence + current_token_choice
373
  )
374
- if current_token_choice_ix == 340:
375
- print("Found info:")
376
- print(f"We generate token '{current_token_choice}', and the total sequence is '{current_sequence}'")
377
  beam_trees[source_beam_ix].children[current_token_choice_ix] = BeamNode(
378
  current_token_ix=current_token_choice_ix,
379
  table=None,
380
  children={},
381
  current_sequence=current_sequence,
382
  cumulative_score=cumulative_score,
383
- total_score=cumulative_score
384
- / ((input_length + step - 1) ** length_penalty),
385
- children_score_divider=((input_length + step) ** length_penalty),
386
- is_final=(
387
- step == len(scores) - 1
388
- or current_token_choice_ix == tokenizer.eos_token_id
389
- ),
390
  is_selected_sequence=(
391
  current_sequence.replace("<|endoftext|>", "")
392
  in [el.replace("<|endoftext|>", "") for el in decoded_sequences]
393
  ),
394
  )
395
 
396
-
397
  # Swap all beams by descending cumul score, so that n°1 has the highest cumulative score, and so on
398
  beam_trees = [
399
- beam_trees[int(top_df_selected.iloc[beam_ix]["beam_index"])]
400
- for beam_ix in range(n_beams)
401
  ]
402
 
403
  # Advance all beams by one token
404
- for beam_ix in range(n_beams):
405
- current_token_choice_ix = top_df_selected.iloc[beam_ix]["token_index"]
406
  beam_trees[beam_ix] = beam_trees[beam_ix].children[current_token_choice_ix]
407
 
 
 
408
  return original_tree
409
 
410
  @spaces.GPU
@@ -459,9 +469,9 @@ with gr.Blocks(
459
  ) as demo:
460
  gr.Markdown(
461
  """# <span style='color:var(--primary-500)!important'>Beam Search Visualizer</span>
462
-
463
  Play with the parameters below to understand how beam search decoding works!
464
-
465
  #### <span style='color:var(--primary-500)!important'>Parameters:</span>
466
  - **Sentence to decode from** (`inputs`): the input sequence to your decoder.
467
  - **Number of steps** (`max_new_tokens`): the number of tokens to generate.
@@ -473,20 +483,20 @@ This parameter will not impact the beam search paths, but only influence the cho
473
  )
474
  text = gr.Textbox(
475
  label="Sentence to decode from",
476
- value="Thank you for",
477
  )
478
  with gr.Row():
479
  n_steps = gr.Slider(
480
- label="Number of steps", minimum=1, maximum=10, step=1, value=4
481
  )
482
  n_beams = gr.Slider(
483
- label="Number of beams", minimum=2, maximum=4, step=1, value=3
484
  )
485
  length_penalty = gr.Slider(
486
  label="Length penalty", minimum=-3, maximum=3, step=0.5, value=1
487
  )
488
  num_return_sequences = gr.Slider(
489
- label="Number of return sequences", minimum=1, maximum=3, step=1, value=2
490
  )
491
 
492
  n_beams.change(
@@ -501,4 +511,4 @@ This parameter will not impact the beam search paths, but only influence the cho
501
  outputs=[out_html, out_markdown],
502
  )
503
 
504
- demo.launch()
 
174
  .nonselected-sequence {
175
  background-color: var(--primary-500);
176
  }
 
 
 
177
  """
178
 
179
 
 
238
  def generate_html(start_sentence, original_tree):
239
  html_output = f"""<div class="custom-container">
240
  <div class="tree">
241
+ <ul> <li> <a href='#' id='root'> <span> <b>{start_sentence}</b> </span> {original_tree.table} </a>"""
242
  html_output += "<ul> "
243
  for subnode in original_tree.children.values():
244
  html_output += generate_nodes(subnode, step=1)
 
270
 
271
 
272
  def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, beam_indexes_source):
 
273
  original_tree = BeamNode(
274
  cumulative_score=0,
275
  current_token_ix=None,
276
  table=None,
277
  current_sequence=start_sentence,
278
  children={},
279
+ children_score_divider=(1 ** length_penalty),
280
  total_score=None,
281
  is_final=False,
282
  is_selected_sequence=False,
283
  )
284
  n_beams = len(scores[0])
285
  beam_trees = [original_tree] * n_beams
286
+ generation_length = len(scores)
287
 
288
  for step, step_scores in enumerate(scores):
289
 
 
294
  beam_indexes,
295
  current_sequence,
296
  top_tokens,
297
+ token_scores,
298
+ ) = ([], [], [], [], [], [])
299
+
300
+ score_idx = 0
301
+ for beam_ix in range(len(beam_trees)):
302
  current_beam = beam_trees[beam_ix]
303
 
304
  # skip if the beam is already final
 
307
 
308
  # Get top cumulative scores for the current beam
309
  current_top_token_indexes = list(
310
+ np.array(scores[step][score_idx].argsort()[-n_beams:])[::-1]
311
  )
312
  top_token_indexes += current_top_token_indexes
313
+ token_scores += list(np.array(scores[step][score_idx][current_top_token_indexes]))
314
  top_cumulative_scores += list(
315
+ np.array(scores[step][score_idx][current_top_token_indexes])
316
  + current_beam.cumulative_score
317
  )
318
  beam_indexes += [beam_ix] * n_beams
319
  current_sequence += [beam_trees[beam_ix].current_sequence] * n_beams
320
  top_tokens += [tokenizer.decode([el]) for el in current_top_token_indexes]
321
+ score_idx += 1
322
 
323
  top_df = pd.DataFrame.from_dict(
324
  {
 
327
  "beam_index": beam_indexes,
328
  "current_sequence": current_sequence,
329
  "token": top_tokens,
330
+ "token_score": token_scores,
331
  }
332
  )
333
  maxes = top_df.groupby(["token_index", "current_sequence"])[
 
336
 
337
  top_df = top_df.loc[maxes]
338
 
339
+ # Sort all top probabilities and keep top n_beams * 2 (* 2 because each beam may end this iteration, and we
340
+ # want to keep at least `n_beams` beams alive)
341
  top_df_selected = top_df.sort_values("cumulative_score", ascending=False).iloc[
342
+ :n_beams * 2
343
  ]
344
+ beams_to_keep = 0
345
+ unfinished_beams = 0
346
+ for _, row in top_df_selected.iterrows():
347
+ beams_to_keep += 1
348
+ current_token_choice_ix = row["token_index"]
349
+ is_final = step == len(scores) - 1 or current_token_choice_ix == tokenizer.eos_token_id
350
+ if not is_final:
351
+ unfinished_beams += 1
352
+ if unfinished_beams >= n_beams:
353
+ break
354
+ if step == generation_length - 1 and beams_to_keep == n_beams:
355
+ break
356
+ top_df_selected_filtered = top_df_selected.iloc[:beams_to_keep]
357
 
358
  # Write the scores table - one per beam source
359
+ score_idx = 0
360
+ for beam_ix in range(len(beam_trees)):
361
  current_beam = beam_trees[beam_ix]
362
  if current_beam.table is None:
363
+ selected_tokens = top_df_selected_filtered.loc[
364
+ top_df_selected_filtered["current_sequence"] == current_beam.current_sequence
365
  ]
366
  markdown_table = generate_markdown_table(
367
+ step_scores[score_idx, :],
368
  current_beam.cumulative_score,
369
  current_beam.children_score_divider,
370
  chosen_tokens=list(selected_tokens["token"].values),
371
  )
372
  beam_trees[beam_ix].table = markdown_table
373
+ if not current_beam.is_final:
374
+ score_idx = min(score_idx + 1, n_beams - 1)
375
 
376
  # Add new children to each beam
377
  cumulative_scores = [beam.cumulative_score for beam in beam_trees]
378
+ for _, row in top_df_selected_filtered.iterrows():
379
  # Update the source tree
380
  source_beam_ix = int(row["beam_index"])
381
  current_token_choice_ix = row["token_index"]
382
  current_token_choice = tokenizer.decode([current_token_choice_ix])
383
+ token_scores = row["token_score"]
384
 
385
+ cumulative_score = cumulative_scores[source_beam_ix] + np.asarray(token_scores)
 
 
 
386
  current_sequence = (
387
  beam_trees[source_beam_ix].current_sequence + current_token_choice
388
  )
389
+ is_final = step == len(scores) - 1 or current_token_choice_ix == tokenizer.eos_token_id
 
 
390
  beam_trees[source_beam_ix].children[current_token_choice_ix] = BeamNode(
391
  current_token_ix=current_token_choice_ix,
392
  table=None,
393
  children={},
394
  current_sequence=current_sequence,
395
  cumulative_score=cumulative_score,
396
+ total_score=cumulative_score / (step + 1 ** length_penalty),
397
+ children_score_divider=((step + 2) ** length_penalty),
398
+ is_final=is_final,
 
 
 
 
399
  is_selected_sequence=(
400
  current_sequence.replace("<|endoftext|>", "")
401
  in [el.replace("<|endoftext|>", "") for el in decoded_sequences]
402
  ),
403
  )
404
 
 
405
  # Swap all beams by descending cumul score, so that n°1 has the highest cumulative score, and so on
406
  beam_trees = [
407
+ beam_trees[int(top_df_selected_filtered.iloc[beam_ix]["beam_index"])]
408
+ for beam_ix in range(beams_to_keep)
409
  ]
410
 
411
  # Advance all beams by one token
412
+ for beam_ix in range(beams_to_keep):
413
+ current_token_choice_ix = top_df_selected_filtered.iloc[beam_ix]["token_index"]
414
  beam_trees[beam_ix] = beam_trees[beam_ix].children[current_token_choice_ix]
415
 
416
+ print(f"Step {step}, beams kept: {beams_to_keep}")
417
+
418
  return original_tree
419
 
420
  @spaces.GPU
 
469
  ) as demo:
470
  gr.Markdown(
471
  """# <span style='color:var(--primary-500)!important'>Beam Search Visualizer</span>
472
+
473
  Play with the parameters below to understand how beam search decoding works!
474
+
475
  #### <span style='color:var(--primary-500)!important'>Parameters:</span>
476
  - **Sentence to decode from** (`inputs`): the input sequence to your decoder.
477
  - **Number of steps** (`max_new_tokens`): the number of tokens to generate.
 
483
  )
484
  text = gr.Textbox(
485
  label="Sentence to decode from",
486
+ value="Conclusion: thanks a lot. That's all for today",
487
  )
488
  with gr.Row():
489
  n_steps = gr.Slider(
490
+ label="Number of steps", minimum=1, maximum=10, step=1, value=10
491
  )
492
  n_beams = gr.Slider(
493
+ label="Number of beams", minimum=2, maximum=4, step=1, value=4
494
  )
495
  length_penalty = gr.Slider(
496
  label="Length penalty", minimum=-3, maximum=3, step=0.5, value=1
497
  )
498
  num_return_sequences = gr.Slider(
499
+ label="Number of return sequences", minimum=1, maximum=4, step=1, value=3
500
  )
501
 
502
  n_beams.change(
 
511
  outputs=[out_html, out_markdown],
512
  )
513
 
514
+ demo.launch()