Spaces:
Running
Running
| import gradio as gr | |
| import json | |
| import matplotlib.pyplot as plt | |
| import pandas as pd | |
| import io | |
| import base64 | |
| import math | |
| import ast | |
| import logging | |
| # Set up logging | |
| logging.basicConfig(level=logging.DEBUG) | |
| logger = logging.getLogger(__name__) | |
| # Function to safely parse JSON or Python dictionary input | |
| def parse_input(json_input): | |
| logger.debug("Attempting to parse input: %s", json_input) | |
| try: | |
| # Try to parse as JSON first | |
| data = json.loads(json_input) | |
| logger.debug("Successfully parsed as JSON") | |
| return data | |
| except json.JSONDecodeError as e: | |
| logger.error("JSON parsing failed: %s", str(e)) | |
| try: | |
| # If JSON fails, try to parse as Python literal (e.g., with single quotes) | |
| data = ast.literal_eval(json_input) | |
| logger.debug("Successfully parsed as Python literal") | |
| # Convert Python dictionary to JSON-compatible format (replace single quotes with double quotes) | |
| def dict_to_json(obj): | |
| if isinstance(obj, dict): | |
| return {str(k): dict_to_json(v) for k, v in obj.items()} | |
| elif isinstance(obj, list): | |
| return [dict_to_json(item) for item in obj] | |
| else: | |
| return obj | |
| converted_data = dict_to_json(data) | |
| logger.debug("Converted to JSON-compatible format") | |
| return converted_data | |
| except (SyntaxError, ValueError) as e: | |
| logger.error("Python literal parsing failed: %s", str(e)) | |
| raise ValueError(f"Malformed input: {str(e)}. Ensure property names are in double quotes (e.g., \"content\") or correct Python dictionary format.") | |
| # Function to ensure a value is a float, converting from string if necessary | |
| def ensure_float(value): | |
| if value is None: | |
| return None | |
| if isinstance(value, str): | |
| try: | |
| return float(value) | |
| except ValueError: | |
| logger.error("Failed to convert string '%s' to float", value) | |
| return None | |
| if isinstance(value, (int, float)): | |
| return float(value) | |
| return None | |
| # Function to process and visualize log probs | |
| def visualize_logprobs(json_input): | |
| try: | |
| # Parse the input (handles both JSON and Python dictionaries) | |
| data = parse_input(json_input) | |
| # Ensure data is a list or dictionary with 'content' | |
| if isinstance(data, dict) and "content" in data: | |
| content = data["content"] | |
| elif isinstance(data, list): | |
| content = data | |
| else: | |
| raise ValueError("Input must be a list or dictionary with 'content' key") | |
| # Extract tokens and log probs, skipping None or non-finite values | |
| tokens = [] | |
| logprobs = [] | |
| for entry in content: | |
| logprob = ensure_float(entry.get("logprob", None)) | |
| if logprob is not None and math.isfinite(logprob): | |
| tokens.append(entry["token"]) | |
| logprobs.append(logprob) | |
| else: | |
| logger.debug("Skipping entry with logprob: %s (type: %s)", entry.get("logprob"), type(entry.get("logprob", None))) | |
| # Prepare table data, handling None in top_logprobs | |
| table_data = [] | |
| for entry in content: | |
| logprob = ensure_float(entry.get("logprob", None)) | |
| # Only include entries with finite logprob and non-None top_logprobs | |
| if ( | |
| logprob is not None | |
| and math.isfinite(logprob) | |
| and "top_logprobs" in entry | |
| and entry["top_logprobs"] is not None | |
| ): | |
| token = entry["token"] | |
| logger.debug("Processing token: %s, logprob: %s (type: %s)", token, logprob, type(logprob)) | |
| top_logprobs = entry["top_logprobs"] | |
| # Ensure all values in top_logprobs are floats | |
| finite_top_logprobs = {} | |
| for key, value in top_logprobs.items(): | |
| float_value = ensure_float(value) | |
| if float_value is not None and math.isfinite(float_value): | |
| finite_top_logprobs[key] = float_value | |
| # Extract top 3 alternatives from top_logprobs | |
| top_3 = sorted( | |
| finite_top_logprobs.items(), key=lambda x: x[1], reverse=True | |
| )[:3] | |
| row = [token, f"{logprob:.4f}"] | |
| for alt_token, alt_logprob in top_3: | |
| row.append(f"{alt_token}: {alt_logprob:.4f}") | |
| # Pad with empty strings if fewer than 3 alternatives | |
| while len(row) < 5: | |
| row.append("") | |
| table_data.append(row) | |
| # Create the plot | |
| if logprobs: | |
| plt.figure(figsize=(10, 5)) | |
| plt.plot(range(len(logprobs)), logprobs, marker="o", linestyle="-", color="b") | |
| plt.title("Log Probabilities of Generated Tokens") | |
| plt.xlabel("Token Position") | |
| plt.ylabel("Log Probability") | |
| plt.grid(True) | |
| plt.xticks(range(len(logprobs)), tokens, rotation=45, ha="right") | |
| plt.tight_layout() | |
| # Save plot to a bytes buffer | |
| buf = io.BytesIO() | |
| plt.savefig(buf, format="png", bbox_inches="tight") | |
| buf.seek(0) | |
| plt.close() | |
| # Convert to base64 for Gradio | |
| img_bytes = buf.getvalue() | |
| img_base64 = base64.b64encode(img_bytes).decode("utf-8") | |
| img_html = f'<img src="data:image/png;base64,{img_base64}" style="max-width: 100%; height: auto;">' | |
| else: | |
| img_html = "No finite log probabilities to plot." | |
| # Create DataFrame for the table | |
| df = ( | |
| pd.DataFrame( | |
| table_data, | |
| columns=[ | |
| "Token", | |
| "Log Prob", | |
| "Top 1 Alternative", | |
| "Top 2 Alternative", | |
| "Top 3 Alternative", | |
| ], | |
| ) | |
| if table_data | |
| else None | |
| ) | |
| # Generate colored text | |
| if logprobs: | |
| min_logprob = min(logprobs) | |
| max_logprob = max(logprobs) | |
| if max_logprob == min_logprob: | |
| normalized_probs = [0.5] * len(logprobs) | |
| else: | |
| normalized_probs = [ | |
| (lp - min_logprob) / (max_logprob - min_logprob) for lp in logprobs | |
| ] | |
| colored_text = "" | |
| for i, (token, norm_prob) in enumerate(zip(tokens, normalized_probs)): | |
| r = int(255 * (1 - norm_prob)) # Red for low confidence | |
| g = int(255 * norm_prob) # Green for high confidence | |
| b = 0 | |
| color = f"rgb({r}, {g}, {b})" | |
| colored_text += f'<span style="color: {color}; font-weight: bold;">{token}</span>' | |
| if i < len(tokens) - 1: | |
| colored_text += " " | |
| colored_text_html = f"<p>{colored_text}</p>" | |
| else: | |
| colored_text_html = "No finite log probabilities to display." | |
| return img_html, df, colored_text_html | |
| except Exception as e: | |
| logger.error("Visualization failed: %s", str(e)) | |
| return f"Error: {str(e)}", None, None | |
| # Gradio interface | |
| with gr.Blocks(title="Log Probability Visualizer") as app: | |
| gr.Markdown("# Log Probability Visualizer") | |
| gr.Markdown( | |
| "Paste your JSON or Python dictionary log prob data below to visualize the tokens and their probabilities. Ensure property names are in double quotes (e.g., \"content\") for JSON, or use correct Python dictionary format." | |
| ) | |
| json_input = gr.Textbox( | |
| label="JSON Input", | |
| lines=10, | |
| placeholder="Paste your JSON (e.g., {\"content\": [...]}) or Python dict (e.g., {'content': [...]}) here...", | |
| ) | |
| plot_output = gr.HTML(label="Log Probability Plot") | |
| table_output = gr.Dataframe(label="Token Log Probabilities and Top Alternatives") | |
| text_output = gr.HTML(label="Colored Text (Confidence Visualization)") | |
| btn = gr.Button("Visualize") | |
| btn.click( | |
| fn=visualize_logprobs, | |
| inputs=json_input, | |
| outputs=[plot_output, table_output, text_output], | |
| ) | |
| app.launch() |