ERROR:__main__:An error occurred during visualization.
Traceback (most recent call last):
  File "<ipython-input-14-817ebddba675>", line 305, in visualize
    fig.frames = frames
  File "/usr/local/lib/python3.10/dist-packages/plotly/basedatatypes.py", line 717, in __setattr__
    super(BaseFigure, self).__setattr__(prop, value)
  File "/usr/local/lib/python3.10/dist-packages/plotly/basewidget.py", line 784, in frames
    BaseFigureWidget._display_frames_error()
  File "/usr/local/lib/python3.10/dist-packages/plotly/basewidget.py", line 800, in _display_frames_error
    raise ValueError(msg)
ValueError: 
Frames are not supported by the plotly.graph_objs.FigureWidget class.
Note: Frames are supported by the plotly.graph_objs.Figure class
ERROR:__main__:Failed to run the TokenVisualizer.
ERROR:__main__:
Frames are not supported by the plotly.graph_objs.FigureWidget class.
Note: Frames are supported by the plotly.graph_objs.Figure class
import time
from typing import List, Tuple, Callable, Dict, Any
import plotly.graph_objs as go
from plotly.subplots import make_subplots
import random
import string
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class TokenVisualizer:
    """
    A class to visualize the top k probable tokens from a token_id input using Plotly.
    It generates and visualizes tokens dynamically in a loop.
    """

    def __init__(
        self,
        k: int = 5,
        num_tokens: int = 30,
        update_interval: float = 1.0,
        decoding_strategy: str = "greedy",
    ) -> None:
        """
        Initializes the TokenVisualizer.

        Args:
            k (int): Number of top tokens to display.
            num_tokens (int): Number of tokens to generate and visualize.
            update_interval (float): Time interval between updates in seconds.
            decoding_strategy (str): Decoding strategy to use.
        """
        self.k = k
        self.num_tokens = num_tokens
        self.update_interval = update_interval
        self.decoding_strategy = decoding_strategy.lower()
        self.decoding_functions: Dict[str, Callable[[Any], Tuple[str, float]]] = {
            "greedy": self.greedy_decoding,
            "beam_search": self.beam_search,
            "top_k_sampling": self.top_k_sampling,
            "top_p_sampling": self.top_p_sampling,
            "temperature_sampling": self.temperature_sampling,
            "random_sampling": self.random_sampling,
            "diverse_beam_search": self.diverse_beam_search,
        }
        self.validate_parameters()

    def validate_parameters(self) -> None:
        """Validates the initialization parameters."""
        if self.k <= 0:
            logger.error("Parameter k must be positive.")
            raise ValueError("Parameter k must be positive.")
        if self.num_tokens <= 0:
            logger.error("Parameter num_tokens must be positive.")
            raise ValueError("Parameter num_tokens must be positive.")
        if self.decoding_strategy not in self.decoding_functions:
            logger.error(
                f"Decoding strategy '{self.decoding_strategy}' is not supported."
            )
            raise ValueError(
                f"Decoding strategy '{self.decoding_strategy}' is not supported."
            )

    def generate_mock_probabilities(self) -> Dict[str, float]:
        """
        Generates mock token probabilities for demonstration purposes.

        Returns:
            Dict[str, float]: A dictionary mapping tokens to their probabilities.
        """
        tokens = [
            ''.join(random.choices(string.ascii_lowercase, k=3)) for _ in range(50)
        ]
        probabilities = [random.random() for _ in range(50)]
        total = sum(probabilities)
        normalized_probs = [p / total for p in probabilities]
        token_probs = dict(zip(tokens, normalized_probs))
        return token_probs

    def greedy_decoding(self, token_probs: Dict[str, float]) -> Tuple[str, float]:
        """
        Greedy decoding strategy: selects the token with the highest probability.

        Args:
            token_probs (Dict[str, float]): Token probabilities.

        Returns:
            Tuple[str, float]: Selected token and its probability.
        """
        token, prob = max(token_probs.items(), key=lambda item: item[1])
        return token, prob

    def beam_search(self, token_probs: Dict[str, float]) -> Tuple[str, float]:
        """
        Placeholder for Beam Search decoding strategy.

        Args:
            token_probs (Dict[str, float]): Token probabilities.

        Returns:
            Tuple[str, float]: Selected token and its probability.
        """
        # Implement Beam Search logic here
        return self.greedy_decoding(token_probs)

    def top_k_sampling(self, token_probs: Dict[str, float], top_k: int = 5) -> Tuple[str, float]:
        """
        Top-k Sampling decoding strategy.

        Args:
            token_probs (Dict[str, float]): Token probabilities.
            top_k (int): Number of top tokens to consider.

        Returns:
            Tuple[str, float]: Selected token and its probability.
        """
        sorted_tokens = sorted(token_probs.items(), key=lambda item: item[1], reverse=True)
        top_k_tokens = sorted_tokens[:top_k]
        tokens, probs = zip(*top_k_tokens)
        normalized_probs = [p / sum(probs) for p in probs]
        selected_token = random.choices(tokens, weights=normalized_probs, k=1)[0]
        selected_prob = dict(top_k_tokens)[selected_token]
        return selected_token, selected_prob

    def top_p_sampling(self, token_probs: Dict[str, float], top_p: float = 0.9) -> Tuple[str, float]:
        """
        Top-p (Nucleus) Sampling decoding strategy.

        Args:
            token_probs (Dict[str, float]): Token probabilities.
            top_p (float): Cumulative probability threshold.

        Returns:
            Tuple[str, float]: Selected token and its probability.
        """
        sorted_tokens = sorted(token_probs.items(), key=lambda item: item[1], reverse=True)
        cumulative_prob = 0.0
        top_p_tokens = []
        for token, prob in sorted_tokens:
            cumulative_prob += prob
            top_p_tokens.append((token, prob))
            if cumulative_prob >= top_p:
                break
        tokens, probs = zip(*top_p_tokens)
        normalized_probs = [p / sum(probs) for p in probs]
        selected_token = random.choices(tokens, weights=normalized_probs, k=1)[0]
        selected_prob = dict(top_p_tokens)[selected_token]
        return selected_token, selected_prob

    def temperature_sampling(
        self, token_probs: Dict[str, float], temperature: float = 1.0
    ) -> Tuple[str, float]:
        """
        Temperature Sampling decoding strategy.

        Args:
            token_probs (Dict[str, float]): Token probabilities.
            temperature (float): Temperature parameter.

        Returns:
            Tuple[str, float]: Selected token and its probability.
        """
        if temperature <= 0:
            logger.error("Temperature must be greater than 0.")
            raise ValueError("Temperature must be greater than 0.")
        tokens, probs = zip(*token_probs.items())
        adjusted_probs = [p ** (1 / temperature) for p in probs]
        total = sum(adjusted_probs)
        normalized_probs = [p / total for p in adjusted_probs]
        selected_token = random.choices(tokens, weights=normalized_probs, k=1)[0]
        selected_prob = token_probs[selected_token]
        return selected_token, selected_prob

    def random_sampling(self, token_probs: Dict[str, float]) -> Tuple[str, float]:
        """
        Random Sampling decoding strategy.

        Args:
            token_probs (Dict[str, float]): Token probabilities.

        Returns:
            Tuple[str, float]: Selected token and its probability.
        """
        tokens, probs = zip(*token_probs.items())
        selected_token = random.choices(tokens, weights=probs, k=1)[0]
        selected_prob = token_probs[selected_token]
        return selected_token, selected_prob

    def diverse_beam_search(self, token_probs: Dict[str, float]) -> Tuple[str, float]:
        """
        Placeholder for Diverse Beam Search decoding strategy.

        Args:
            token_probs (Dict[str, float]): Token probabilities.

        Returns:
            Tuple[str, float]: Selected token and its probability.
        """
        # Implement Diverse Beam Search logic here
        return self.greedy_decoding(token_probs)

    def get_next_token(self, token_probs: Dict[str, float]) -> Tuple[str, float]:
        """
        Selects the next token based on the decoding strategy.

        Args:
            token_probs (Dict[str, float]): Token probabilities.

        Returns:
            Tuple[str, float]: Selected token and its probability.
        """
        decoding_func = self.decoding_functions[self.decoding_strategy]
        try:
            if self.decoding_strategy == "top_k_sampling":
                return decoding_func(token_probs, top_k=self.k)
            elif self.decoding_strategy == "top_p_sampling":
                return decoding_func(token_probs, top_p=0.9)
            elif self.decoding_strategy == "temperature_sampling":
                return decoding_func(token_probs, temperature=1.0)
            else:
                return decoding_func(token_probs)
        except Exception as e:
            logger.exception("Error in decoding strategy.")
            raise e

    def visualize(self) -> None:
        """
        Generates and visualizes tokens dynamically using Plotly with Play/Pause controls.
        """
        try:
            # Initialize the Plotly figure with initial data
            fig = go.FigureWidget()

            # Initial empty bar
            bar = go.Bar(x=[], y=[], marker_color=[])
            fig.add_trace(bar)

            # Set up the layout with buttons
            fig.update_layout(
                title=f"Top {self.k} Probable Tokens",
                xaxis_title="Tokens",
                yaxis_title="Probability",
                showlegend=False,
                updatemenus=[
                    {
                        "type": "buttons",
                        "buttons": [
                            {
                                "label": "Play",
                                "method": "animate",
                                "args": [None, {"frame": {"duration": self.update_interval * 1000, "redraw": True},
                                               "fromcurrent": True, "transition": {"duration": 0}}],
                            },
                            {
                                "label": "Pause",
                                "method": "animate",
                                "args": [[None], {"frame": {"duration": 0, "redraw": False},
                                                 "mode": "immediate",
                                                 "transition": {"duration": 0}}],
                            },
                        ],
                        "direction": "left",
                        "pad": {"r": 10, "t": 87},
                        "showactive": False,
                        "x": 0.1,
                        "xanchor": "right",
                        "y": 0,
                        "yanchor": "top"
                    }
                ]
            )

            frames = []
            generated_tokens: List[str] = []

            for i in range(self.num_tokens):
                # Simulate fetching token probabilities
                token_probs = self.generate_mock_probabilities()

                # Get top k tokens
                sorted_tokens = sorted(
                    token_probs.items(), key=lambda item: item[1], reverse=True
                )
                top_k = sorted_tokens[: self.k]
                tokens, probs = zip(*top_k)

                # Define colors
                colors = ["red"] + ["blue"] * (self.k - 1)

                # Append the selected token
                selected_token, selected_prob = self.get_next_token(token_probs)
                generated_tokens.append(selected_token)
                logger.info(f"Generated Token {i + 1}: {selected_token} (P={selected_prob:.4f})")

                # Create a frame for the current state
                frame = go.Frame(
                    data=[go.Bar(x=tokens, y=probs, marker_color=colors)],
                    layout=go.Layout(title_text=f"Top {self.k} Probable Tokens - Token {i + 1}")
                )
                frames.append(frame)

            # Assign frames to the figure
            fig.frames = frames

            # Optionally, set the initial data to the first frame
            if frames:
                fig.add_trace(frames[0].data[0])
                fig.layout.title.text = frames[0].layout.title.text

            fig.show()

        except Exception as e:
            logger.exception("An error occurred during visualization.")
            raise e

def main() -> None:
    """
    Main function to run the TokenVisualizer.
    """
    try:
        visualizer = TokenVisualizer(
            k=5,
            num_tokens=30,
            update_interval=0.5,
            decoding_strategy="top_k_sampling",
        )
        visualizer.visualize()
    except Exception as e:
        logger.error("Failed to run the TokenVisualizer.")
        logger.error(str(e))

if __name__ == "__main__":
    main()