ERROR:__main__:An error occurred during visualization.
Traceback (most recent call last):
  File "<ipython-input-12-7761cb4412fc>", line 260, in visualize
    bar.x = tokens
  File "/usr/local/lib/python3.10/dist-packages/plotly/basedatatypes.py", line 720, in __setattr__
    raise AttributeError(prop)
AttributeError: x
ERROR:__main__:Failed to run the TokenVisualizer.
ERROR:__main__:x

we concept play button
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.
        """
        try:
            # Initialize the Plotly figure
            fig = go.FigureWidget()
            fig.update_layout(
                title=f"Top {self.k} Probable Tokens",
                xaxis_title="Tokens",
                yaxis_title="Probability",
                showlegend=False,
            )
            bar = fig.add_bar(x=[], y=[], marker_color=[])
            fig.show()

            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)

                # Update the bar chart
                with fig.batch_update():
                    bar.x = tokens
                    bar.y = probs
                    colors = ["red"] + ["blue"] * (self.k - 1)
                    bar.marker.color = colors

                    fig.layout.title.text = (
                        f"Top {self.k} Probable Tokens - Token {i + 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})")

                # Wait for the specified interval before next update
                time.sleep(self.update_interval)

        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()