"""
for token_idx in np.array(np.argsort(scores)[-top_k:])[::-1]:
token = tokenizer.decode([token_idx])
item_class = ""
if chosen_tokens and token in chosen_tokens:
item_class = "chosen-token"
markdown_table += f"""
"""
return markdown_table
def generate_nodes(node, step):
"""Recursively generate HTML for the tree nodes."""
token = tokenizer.decode([node.current_token_ix])
if node.is_final:
if node.is_selected_sequence:
selected_class = "selected-sequence"
else:
selected_class = "nonselected-sequence"
return f"
"
for subnode in original_tree.children.values():
html_output += generate_nodes(subnode, step=1)
html_output += "
"
html_output += """
"""
return html_output
import pandas as pd
from typing import Dict
from dataclasses import dataclass
@dataclass
class BeamNode:
current_token_ix: int
cumulative_score: float
children_score_divider: float
table: str
current_sequence: str
children: Dict[int, "BeamNode"]
total_score: float
is_final: bool
is_selected_sequence: bool
def generate_beams(start_sentence, scores, length_penalty, decoded_sequences, beam_indexes_source):
original_tree = BeamNode(
cumulative_score=0,
current_token_ix=None,
table=None,
current_sequence=start_sentence,
children={},
children_score_divider=(1 ** length_penalty),
total_score=None,
is_final=False,
is_selected_sequence=False,
)
n_beams = len(scores[0])
beam_trees = [original_tree] * n_beams
generation_length = len(scores)
for step, step_scores in enumerate(scores):
# Gather all possible descendants for each beam
(
top_token_indexes,
top_cumulative_scores,
beam_indexes,
current_sequence,
top_tokens,
token_scores,
) = ([], [], [], [], [], [])
for beam_ix in range(len(beam_trees)): # not range(n_beams) since there might be more ongoing trees.
current_beam = beam_trees[beam_ix]
# skip if the beam is already final
if current_beam.is_final:
continue
# Get top cumulative scores for the current beam
current_top_token_indexes = list(
np.array(scores[step][beam_ix].argsort()[-n_beams:])[::-1]
)
top_token_indexes += current_top_token_indexes
token_scores += list(np.array(scores[step][beam_ix][current_top_token_indexes]))
top_cumulative_scores += list(
np.array(scores[step][beam_ix][current_top_token_indexes])
+ current_beam.cumulative_score
)
beam_indexes += [beam_ix] * n_beams
current_sequence += [beam_trees[beam_ix].current_sequence] * n_beams
top_tokens += [tokenizer.decode([el]) for el in current_top_token_indexes]
top_df = pd.DataFrame.from_dict(
{
"token_index": top_token_indexes,
"cumulative_score": top_cumulative_scores,
"beam_index": beam_indexes,
"current_sequence": current_sequence,
"token": top_tokens,
"token_score": token_scores,
}
)
maxes = top_df.groupby(["token_index", "current_sequence"])[
"cumulative_score"
].idxmax()
top_df = top_df.loc[maxes]
# Sort all top probabilities and keep top n_beams * 2 (* 2 because each beam may end this iteration, and we
# want to keep at least `n_beams` beams alive)
top_df_selected = top_df.sort_values("cumulative_score", ascending=False).iloc[
:n_beams * 2
]
beams_to_keep = 0
unfinished_beams = 0
for _, row in top_df_selected.iterrows():
beams_to_keep += 1
current_token_choice_ix = row["token_index"]
is_final = step == len(scores) - 1 or current_token_choice_ix == tokenizer.eos_token_id
if not is_final:
unfinished_beams += 1
if unfinished_beams >= n_beams:
break
if step == generation_length - 1 and beams_to_keep == n_beams:
break
top_df_selected_filtered = top_df_selected.iloc[:beams_to_keep]
# Write the scores table in each beam tree
score_idx = 0
for beam_ix in range(len(beam_trees)):
current_beam = beam_trees[beam_ix]
if current_beam.table is None:
selected_tokens = top_df_selected_filtered.loc[
top_df_selected_filtered["current_sequence"] == current_beam.current_sequence
]
markdown_table = generate_markdown_table(
step_scores[score_idx, :],
current_beam.cumulative_score,
current_beam.children_score_divider,
chosen_tokens=list(selected_tokens["token"].values),
)
beam_trees[beam_ix].table = markdown_table
if not current_beam.is_final:
score_idx = min(score_idx + 1, n_beams - 1)
# Add new children to each beam
cumulative_scores = [beam.cumulative_score for beam in beam_trees]
for _, row in top_df_selected_filtered.iterrows():
# Update the source tree
source_beam_ix = int(row["beam_index"])
current_token_choice_ix = row["token_index"]
current_token_choice = tokenizer.decode([current_token_choice_ix])
token_scores = row["token_score"]
cumulative_score = cumulative_scores[source_beam_ix] + np.asarray(token_scores)
current_sequence = (
beam_trees[source_beam_ix].current_sequence + current_token_choice
)
is_final = step == len(scores) - 1 or current_token_choice_ix == tokenizer.eos_token_id
beam_trees[source_beam_ix].children[current_token_choice_ix] = BeamNode(
current_token_ix=current_token_choice_ix,
table=None,
children={},
current_sequence=current_sequence,
cumulative_score=cumulative_score,
total_score=cumulative_score / (step + 1 ** length_penalty),
children_score_divider=((step + 2) ** length_penalty),
is_final=is_final,
is_selected_sequence=(
current_sequence.replace("<|endoftext|>", "")
in [el.replace("<|endoftext|>", "") for el in decoded_sequences]
),
)
# Swap all beams by descending cumul score, so that n°1 has the highest cumulative score, and so on
beam_trees = [
beam_trees[int(top_df_selected_filtered.iloc[beam_ix]["beam_index"])]
for beam_ix in range(beams_to_keep)
]
# Advance all beams by one token
for beam_ix in range(beams_to_keep):
current_token_choice_ix = top_df_selected_filtered.iloc[beam_ix]["token_index"]
beam_trees[beam_ix] = beam_trees[beam_ix].children[current_token_choice_ix]
print(f"Step {step}, beams kept: {beams_to_keep}")
return original_tree
@spaces.GPU
def get_beam_search_html(
input_text, number_steps, number_beams, length_penalty, num_return_sequences
):
inputs = tokenizer([input_text], return_tensors="pt")
outputs = model.generate(
**inputs,
max_new_tokens=number_steps,
num_beams=number_beams,
num_return_sequences=num_return_sequences,
return_dict_in_generate=True,
length_penalty=length_penalty,
output_scores=True,
do_sample=False,
)
markdown = "The conclusive sequences are the ones that end in an `<|endoftext|>` token or at the end of generation."
markdown += "\n\nThey are ranked by their scores, as given by the formula `score = cumulative_score / (output_length ** length_penalty)`.\n\n"
markdown += "Only the top `num_beams` scoring sequences are returned: in the tree they are highlighted in **