build-small-hackathon/InflectionLM
1
1import gradio as gr2import spaces3import math4from inflections_funcs import start_model, make_beams, get_beam_tokens, calculate_score_vectors5 6# Global model initialization7print("Loading model...")8model, processor = start_model()9print("Model loaded.")10 11def generate_beam_html(index, all_beams_data, dark_mode):12 """13 Helper to construct the main and detailed HTML for a specific beam index.14 """15 beam_tokens = all_beams_data["beam_tokens"][index]16 beam_scores = all_beams_data["score_vectors"][index]17 18 # Use log-prob (beam search) or average prob (sampling)19 raw_score = all_beams_data["sequences_scores"][index]20 if all_beams_data.get("is_sampling", False):21 beam_overall_score = raw_score22 else:23 beam_overall_score = math.exp(raw_score)24 25 # Construct Main HTML output26 main_container_style = "border: 2px solid lightblue; padding: 15px; border-radius: 8px; font-family: sans-serif;"27 title_style = "font-weight: bold; margin-bottom: 10px; font-size: 1.1em;"28 29 main_html = f'<div style="{main_container_style}">'30 main_html += f'<div style="{title_style}">Response {index + 1} | Score: {beam_overall_score:.4f}</div>'31 main_html += '<div style="font-family: monospace; white-space: pre-wrap; line-height: 1.5; overflow-wrap: break-word; word-wrap: break-word;">'32 33 normal_color = "white" if dark_mode else "black"34 35 for token, score in zip(beam_tokens, beam_scores):36 # Comprehensive replacement of SentencePiece space (U+2581),37 # literal underscores (U+005F), and non-breaking spaces (U+00A0).38 display_token = token.replace('▁', ' ').replace('_', ' ').replace(' ', ' ')39 if score < 0.6:40 color = "red"41 bg_color = "#441111" if dark_mode else "#ffe6e6"42 style = f"color: {color}; background-color: {bg_color};"43 else:44 color = normal_color45 style = f"color: {color};"46 main_html += f'<span style="{style}" title="Score: {score:.4f}">{display_token}</span>'47 main_html += '</div></div>'48 49 # Construct Detailed HTML output50 detail_table_style = "border-collapse: collapse; width: 100%; max-width: 500px; font-family: monospace;"51 th_style = "border: 1px solid #ccc; padding: 8px; text-align: left; background-color: #f2f2f2;" if not dark_mode else "border: 1px solid #444; padding: 8px; text-align: left; background-color: #333;"52 td_style = "border: 1px solid #ccc; padding: 8px; text-align: left;" if not dark_mode else "border: 1px solid #444; padding: 8px; text-align: left;"53 54 detail_html = f'<table style="{detail_table_style}"><thead><tr><th style="{th_style}">Token</th><th style="{th_style}">Score</th></tr></thead><tbody>'55 for token, score in zip(beam_tokens, beam_scores):56 display_token = token.replace(' ', ' ').replace('_', ' ').replace(' ', ' ')57 if score < 0.6:58 color = "red"59 bg_color = "#441111" if dark_mode else "#ffe6e6"60 token_cell_style = f"{td_style} color: {color}; background-color: {bg_color};"61 else:62 color = normal_color63 token_cell_style = f"{td_style} color: {color};"64 detail_html += f'<tr><td style="{token_cell_style}">{display_token}</td><td style="{td_style}">{score:.4f}</td></tr>'65 detail_html += '</tbody></table>'66 67 return main_html, detail_html68 69@spaces.GPU70def predict(prompt, dark_mode, temperature):71 """72 Generates responses for 3 beams and returns the first beam's visualization and visibility for controls.73 """74 # Generate beams75 generated_dicts, transcription = make_beams(model, processor, prompt, temperature=temperature)76 77 # Get tokens and scores for all beams78 beam_tokens = get_beam_tokens(generated_dicts, processor)79 score_vectors = calculate_score_vectors(model, generated_dicts)80 81 if not beam_tokens or not score_vectors:82 return "<span style='color: grey'>No tokens generated.</span>", "", gr.update(visible=False), gr.update(visible=False), None, 083 84 # Initialize state with all beam data85 # Convert tensors to lists to avoid ZeroGPU serialization errors86 87 # Safely handle sequence scores88 if hasattr(generated_dicts, 'sequences_scores') and generated_dicts.sequences_scores is not None:89 seq_scores = generated_dicts.sequences_scores.tolist() if hasattr(generated_dicts.sequences_scores, 'tolist') else generated_dicts.sequences_scores90 is_sampling = False91 else:92 # For sampling, approximate overall score as the average probability of tokens in the beam93 seq_scores = [sum(scores) / len(scores) if scores else 0.0 for scores in score_vectors]94 is_sampling = True95 96 all_beams_data = {97 "beam_tokens": beam_tokens,98 "score_vectors": score_vectors,99 "sequences_scores": seq_scores,100 "is_sampling": is_sampling101 }102 103 # Generate HTML for the first beam (index 0)104 main_html, detail_html = generate_beam_html(0, all_beams_data, dark_mode)105 106 return main_html, detail_html, gr.update(visible=True), gr.update(visible=True), all_beams_data, 0107 108def switch_beam(current_index, all_beams_data, dark_mode):109 """110 Increments the beam index and returns the updated HTML.111 """112 if all_beams_data is None:113 return None, None, 0114 115 new_index = (current_index + 1) % 3116 main_html, detail_html = generate_beam_html(new_index, all_beams_data, dark_mode)117 return main_html, detail_html, new_index118 119# Gradio Interface120with gr.Blocks() as demo:121 with gr.Row():122 with gr.Column(scale=1):123 gr.Image("Stochastic_parrot.JPG")124 with gr.Column(scale=3):125 gr.Markdown("# InflectionLM: Output and Token Visualization")126 gr.Markdown('''Input a prompt. The model will:127 - Generate multiple independent responses using Top-P and Top-K sampling to ensure diversity.128 - Display the first response, highlighting any word or parts of words (tokens) with a probability score < 0.6 in red.129 - These unconfident words/tokens can be considered inflection points, or places where the output could change easily.130 - Provide a toggle button to view detailed word/token probabilities for the response.131 - Switch between the generated responses to see different possible outputs and their associated token probabilities.132 ''')133 134 with gr.Column():135 prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=1, show_label=True)136 with gr.Row():137 dark_mode_toggle = gr.Checkbox(label="Dark Mode", value=True)138 temp_radio = gr.Radio(139 label="Temperature",140 choices=[0.0, 0.5, 1.0, 1.5],141 value=1.0142 )143 submit_btn = gr.Button("Generate")144 145 output_html = gr.HTML(label="Highlighted Beam")146 147 with gr.Row():148 toggle_btn = gr.Button("Show Token Details", visible=False)149 next_beam_btn = gr.Button("Next Beam", visible=False)150 151 detail_html = gr.HTML(label="Detailed Token Probabilities", visible=False)152 153 # State components154 visibility_state = gr.State(value=False)155 all_beams_state = gr.State(value=None)156 current_beam_state = gr.State(value=0)157 158 def toggle_view(visible):159 return not visible, gr.update(visible=not visible)160 161 toggle_btn.click(fn=toggle_view, inputs=visibility_state, outputs=[visibility_state, detail_html])162 163 # Switch beam logic164 next_beam_btn.click(165 fn=switch_beam,166 inputs=[current_beam_state, all_beams_state, dark_mode_toggle],167 outputs=[output_html, detail_html, current_beam_state]168 )169 170 # Trigger generation on both button click and Enter key in textbox171 submit_btn.click(172 fn=predict,173 inputs=[prompt_input, dark_mode_toggle, temp_radio],174 outputs=[output_html, detail_html, toggle_btn, next_beam_btn, all_beams_state, current_beam_state]175 )176 prompt_input.submit(177 fn=predict,178 inputs=[prompt_input, dark_mode_toggle, temp_radio],179 outputs=[output_html, detail_html, toggle_btn, next_beam_btn, all_beams_state, current_beam_state]180 )181 182if __name__ == "__main__":183 demo.launch()184