beam_search_visualizer

Running on T4

App Files Files Community

m-ric commited on Mar 21, 2024

Commit

a7a4e14

verified ·

1 Parent(s): b20c42a

Update app.py

Browse files

Files changed (1) hide show

app.py +14 -63

app.py CHANGED Viewed

@@ -1,22 +1,15 @@
-import gradio as gr
 STYLE = """
-@import url('https://fonts.googleapis.com/css2?family=Poppins:ital,wght@0,100;0,200;0,300;0,400;0,500;0,600;0,700;0,800;0,900;1,100;1,200;1,300;1,400;1,500;1,600;1,700;1,800;1,900&display=swap');
-* {
-	padding: 0px;
-	margin: 0px;
-	box-sizing: border-box;
-    font-size: 16px;
-}
 body {
 	height: 100vh;
 	width: 100vw;
 	display: grid;
 	align-items: center;
-	font-family: 'Poppins', sans-serif;
 }
 .tree {
 	width: 100%;
 	height: auto;
 	text-align: center;
@@ -27,8 +20,7 @@ body {
 	transition: .5s;
 }
 .tree li {
-	display: flex;
-    flex-direction:row;
 	text-align: center;
 	list-style-type: none;
 	position: relative;
@@ -87,13 +79,6 @@ body {
 	border-radius: 5px;
 	transition: .5s;
 }
-.tree li a img {
-	width: 50px;
-	height: 50px;
-	margin-bottom: 10px !important;
-	border-radius: 100px;
-	margin: auto;
-}
 .tree li a span {
 	border: 1px solid #ccc;
 	border-radius: 5px;
@@ -122,56 +107,19 @@ tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 model = AutoModelForCausalLM.from_pretrained("gpt2")
 tokenizer.pad_token_id = tokenizer.eos_token_id
-def display_top_k_tokens(scores, sequences, beam_indices):
-    display = "<div style='display: flex; flex-direction:row;'>"
-    for i, sequence in enumerate(sequences):
-        markdown_table = f"""<p>Sequence {i}: {tokenizer.batch_decode(sequence)}<p><br>
-        <table>
-            <tr>
-                <th><b>Token</b></th>
-                <th><b>Probability</b></th>
-            </tr>"""
-        for step, step_scores in enumerate(scores):
-            markdown_table += f"""
-            <tr>
-                <td><b>Step {step}</b></td>
-                <td>=====</td>
-            </tr>"""
-            current_beam = beam_indices[i, step]
-            chosen_token = sequences[i, step]
-            for token_idx in np.argsort(step_scores[current_beam, :])[-5:]:
-                if token_idx == chosen_token:
-                    markdown_table += f"""
-                        <tr style="background-color:red">
-                            <td>{tokenizer.decode([token_idx])}</td>
-                            <td>{step_scores[current_beam, token_idx]}</td>
-                        </tr>"""
-                else:
-                    markdown_table += f"""
-                        <tr>
-                            <td>{tokenizer.decode([token_idx])}</td>
-                            <td>{step_scores[current_beam, token_idx]}</td>
-                        </tr>"""
-        markdown_table += "</table>"
-        display += markdown_table
-    display += "</div>"
-    print(display)
-    return display
 def generate_html(token, node):
     """Recursively generate HTML for the tree."""
-    html_content = f" <ul> <a href='#'> <p> <b>{token}</b> </p> "
     html_content += node["table"] if node["table"] is not None else ""
     html_content += "</a>"
     if len(node["children"].keys()) > 0:
-        html_content += "<li> "
         for token, subnode in node["children"].items():
             html_content += generate_html(token, subnode)
-        html_content += "</li>"
-    html_content += "</ul>"
     return html_content
@@ -202,7 +150,8 @@ def display_tree(scores, sequences, beam_indices):
     display = """<body>
 		<div class="container">
 			<div class="row">
-				<div class="tree">"""
     sequences = sequences.cpu().numpy()
     print(tokenizer.batch_decode(sequences))
     original_tree = {"table": None, "children": {}}
@@ -230,6 +179,7 @@ def display_tree(scores, sequences, beam_indices):
     display += generate_html("Today is", original_tree)
     display += """
         </div>
     </div>
     </div>
@@ -260,7 +210,8 @@ def get_tables(input_text, number_steps, number_beams):
         outputs.beam_indices[:, : -len(inputs)],
     )
     return tables
 with gr.Blocks(
     theme=gr.themes.Soft(

 STYLE = """
 body {
 	height: 100vh;
 	width: 100vw;
 	display: grid;
 	align-items: center;
 }
 .tree {
+	padding: 0px;
+	margin: 0px;
+	box-sizing: border-box;
+    font-size: 16px;
 	width: 100%;
 	height: auto;
 	text-align: center;
 	transition: .5s;
 }
 .tree li {
+	display: inline-table;
 	text-align: center;
 	list-style-type: none;
 	position: relative;
 	border-radius: 5px;
 	transition: .5s;
 }
 .tree li a span {
 	border: 1px solid #ccc;
 	border-radius: 5px;
 model = AutoModelForCausalLM.from_pretrained("gpt2")
 tokenizer.pad_token_id = tokenizer.eos_token_id
 def generate_html(token, node):
     """Recursively generate HTML for the tree."""
+    html_content = f" <li> <a href='#'> <span> <b>{token}</b> </span> "
     html_content += node["table"] if node["table"] is not None else ""
     html_content += "</a>"
     if len(node["children"].keys()) > 0:
+        html_content += "<ul> "
         for token, subnode in node["children"].items():
             html_content += generate_html(token, subnode)
+        html_content += "</ul>"
+    html_content += "</li>"
     return html_content
     display = """<body>
 		<div class="container">
 			<div class="row">
+				<div class="tree">
+                <ul>"""
     sequences = sequences.cpu().numpy()
     print(tokenizer.batch_decode(sequences))
     original_tree = {"table": None, "children": {}}
     display += generate_html("Today is", original_tree)
     display += """
+        </ul>
         </div>
     </div>
     </div>
         outputs.beam_indices[:, : -len(inputs)],
     )
     return tables
+import gradio as gr
 with gr.Blocks(
     theme=gr.themes.Soft(