4notes/LLMs/distributiontorchmodule.py
C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\gradio\utils.py:1002: UserWarning: Expected 3 arguments for function <function plot_function at 0x0000020745719260>, received 5.
  warnings.warn(
C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\gradio\utils.py:1010: UserWarning: Expected maximum 3 arguments for function <function plot_function at 0x0000020745719260>, received 5.
  warnings.warn(
Running on local URL:  http://127.0.0.1:7860

To create a public link, set `share=True` in `launch()`.
C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\gradio\analytics.py:106: UserWarning: IMPORTANT: You are using gradio version 4.41.0, however version 4.44.1 is available, please upgrade.
--------
  warnings.warn(
Traceback (most recent call last):
  File "C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\gradio\queueing.py", line 536, in process_events
    response = await route_utils.call_process_api(
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\gradio\route_utils.py", line 288, in call_process_api
    output = await app.get_blocks().process_api(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\gradio\blocks.py", line 1941, in process_api        
    data = await self.postprocess_data(block_fn, result["prediction"], state)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\gradio\blocks.py", line 1759, in postprocess_data   
    raise InvalidComponentError(
gradio.exceptions.InvalidComponentError: <class 'gradio.layouts.column.Column'> Component not a valid output component.
Traceback (most recent call last):
  File "C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\gradio\queueing.py", line 536, in process_events
    response = await route_utils.call_process_api(
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\gradio\route_utils.py", line 288, in call_process_api
    output = await app.get_blocks().process_api(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\gradio\blocks.py", line 1931, in process_api        
    result = await self.call_function(
             ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\gradio\blocks.py", line 1516, in call_function      
    prediction = await anyio.to_thread.run_sync(  # type: ignore
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\anyio\to_thread.py", line 56, in run_sync
    return await get_async_backend().run_sync_in_worker_thread(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\anyio\_backends\_asyncio.py", line 2177, in run_sync_in_worker_thread
    return await future
           ^^^^^^^^^^^^
  File "C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\anyio\_backends\_asyncio.py", line 859, in run      
    result = context.run(func, *args)
             ^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\gradio\utils.py", line 826, in wrapper
    response = f(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^
TypeError: plot_function() takes 3 positional arguments but 5 were given
import torch
import gradio as gr
import plotly.graph_objs as go
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import numpy as np


class DistributionInfo:
    """
    A class to hold information about each distribution.
    """

    def __init__(
        self,
        cls: Callable[..., torch.distributions.Distribution],
        parameters: Dict[str, Dict[str, Any]],
        support: Optional[Tuple[float, float]] = None,
    ) -> None:
        """
        Initializes the DistributionInfo.

        Args:
            cls: The PyTorch distribution class.
            parameters: A dictionary mapping parameter names to their type and default value.
            support: The support of the distribution as a tuple (min, max).
        """
        self.cls = cls
        self.parameters = parameters
        self.support = support


# Mapping of distribution names to their corresponding torch.distributions classes and parameters
DISTRIBUTIONS: Dict[str, DistributionInfo] = {
    "Bernoulli": DistributionInfo(
        cls=torch.distributions.Bernoulli,
        parameters={
            "probs": {"type": "float", "default": 0.5},
        },
        support=(0, 1),
    ),
    "Beta": DistributionInfo(
        cls=torch.distributions.Beta,
        parameters={
            "concentration1": {"type": "float", "default": 1.0},
            "concentration0": {"type": "float", "default": 1.0},
        },
        support=(0, 1),
    ),
    "Binomial": DistributionInfo(
        cls=torch.distributions.Binomial,
        parameters={
            "total_count": {"type": "int", "default": 10},
            "probs": {"type": "float", "default": 0.5},
        },
        support=(0, None),
    ),
    "Categorical": DistributionInfo(
        cls=torch.distributions.Categorical,
        parameters={
            "probs": {"type": "list", "default": [0.25, 0.25, 0.25, 0.25]},
        },
        support=(0, None),
    ),
    "Normal": DistributionInfo(
        cls=torch.distributions.Normal,
        parameters={
            "loc": {"type": "float", "default": 0.0},
            "scale": {"type": "float", "default": 1.0},
        },
        support=(-np.inf, np.inf),
    ),
    # Add more distributions here following the same structure
    # ...
}


def create_distribution(
    distribution_name: str, params: Dict[str, Any]
) -> torch.distributions.Distribution:
    """
    Creates a PyTorch distribution object based on the distribution name and parameters.

    Args:
        distribution_name: The name of the distribution.
        params: A dictionary of parameters for the distribution.

    Returns:
        An instance of a PyTorch distribution.

    Raises:
        ValueError: If the distribution is not supported or parameters are invalid.
    """
    if distribution_name not in DISTRIBUTIONS:
        raise ValueError(f"Unsupported distribution: {distribution_name}")

    dist_info = DISTRIBUTIONS[distribution_name]
    try:
        # Convert list parameters to torch tensors if necessary
        for param, value in params.items():
            if dist_info.parameters[param]["type"] == "list":
                params[param] = torch.tensor(value, dtype=torch.float32)
        distribution = dist_info.cls(**params)
    except Exception as e:
        raise ValueError(f"Error creating distribution: {e}")

    return distribution


def generate_x_values(
    support: Tuple[float, float], num_points: int = 1000
) -> torch.Tensor:
    """
    Generates x values over the support of the distribution.

    Args:
        support: A tuple indicating the range (min, max).
        num_points: Number of points to generate.

    Returns:
        A tensor of x values.
    """
    min_val, max_val = support
    if min_val == -np.inf:
        min_val = -10
    if max_val == np.inf:
        max_val = 10
    return torch.linspace(min_val, max_val, steps=num_points)


def plot_distribution(
    distribution: torch.distributions.Distribution,
    plot_pdf: bool,
    plot_cdf: bool,
) -> Tuple[go.Figure, str]:
    """
    Generates Plotly figures for PDF and/or CDF of the distribution.

    Args:
        distribution: The PyTorch distribution object.
        plot_pdf: Whether to plot PDF.
        plot_cdf: Whether to plot CDF.

    Returns:
        A tuple containing the Plotly figure and an error message if any.
    """
    try:
        # Determine support for x-axis
        if hasattr(distribution, "support"):
            support = distribution.support.lower_bound, distribution.support.upper_bound
        else:
            support = (-10, 10)

        x = generate_x_values(support)
        x_np = x.numpy()

        fig = go.Figure()
        if plot_pdf:
            pdf_values = distribution.log_prob(x).exp().numpy()
            fig.add_trace(
                go.Scatter(
                    x=x_np,
                    y=pdf_values,
                    mode="lines",
                    name="PDF",
                )
            )
        if plot_cdf:
            cdf_values = distribution.cdf(x).numpy()
            fig.add_trace(
                go.Scatter(
                    x=x_np,
                    y=cdf_values,
                    mode="lines",
                    name="CDF",
                )
            )
        fig.update_layout(
            title="Distribution Plot",
            xaxis_title="x",
            yaxis_title="Probability",
            template="plotly_white",
        )
        return fig, ""
    except Exception as e:
        return go.Figure(), f"Error generating plot: {e}"


def validate_and_parse_params(
    distribution_name: str, inputs: Dict[str, Any]
) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
    """
    Validates and parses the input parameters for the selected distribution.

    Args:
        distribution_name: The name of the distribution.
        inputs: A dictionary of input parameters as strings.

    Returns:
        A tuple containing the parsed parameters dictionary and an error message if any.
    """
    dist_info = DISTRIBUTIONS.get(distribution_name)
    if not dist_info:
        return None, f"Unsupported distribution: {distribution_name}"

    parsed_params = {}
    for param, prop in dist_info.parameters.items():
        value = inputs.get(param)
        try:
            if prop["type"] == "float":
                parsed_params[param] = float(value)
            elif prop["type"] == "int":
                parsed_params[param] = int(value)
            elif prop["type"] == "list":
                # Expecting a comma-separated string of numbers
                parsed_params[param] = [float(x.strip()) for x in value.split(",")]
            else:
                parsed_params[param] = value
        except ValueError:
            return None, f"Invalid value for parameter '{param}': {value}"
    return parsed_params, None


def get_parameter_components(
    distribution_name: str,
) -> List[gr.components.Component]:
    """
    Generates a list of Gradio input components based on the distribution's parameters.

    Args:
        distribution_name: The name of the distribution.

    Returns:
        A list of Gradio components corresponding to the distribution's parameters.
    """
    dist_info = DISTRIBUTIONS.get(distribution_name)
    if not dist_info:
        return []

    components = []
    for param, prop in dist_info.parameters.items():
        if prop["type"] == "float":
            components.append(
                gr.Number(
                    label=param,
                    value=prop["default"],
                )
            )
        elif prop["type"] == "int":
            components.append(
                gr.Number(
                    label=param,
                    value=prop["default"],
                    precision=0,
                )
            )
        elif prop["type"] == "list":
            components.append(
                gr.Textbox(
                    label=f"{param} (comma separated)",
                    value=",".join(map(str, prop["default"])),
                )
            )
        else:
            components.append(
                gr.Textbox(
                    label=param,
                    value=str(prop["default"]),
                )
            )
    return components


def update_parameters(distribution_name: str) -> List[gr.components.Component]:
    """
    Updates the parameter input fields based on the selected distribution.

    Args:
        distribution_name: The name of the selected distribution.

    Returns:
        A list of updated Gradio components.
    """
    return get_parameter_components(distribution_name)


def plot_function(
    distribution_name: str,
    plot_pdf: bool,
    plot_cdf: bool,
    **kwargs,
) -> Union[go.Figure, str]:
    """
    The main function to create distribution and generate plots.

    Args:
        distribution_name: The name of the distribution.
        plot_pdf: Whether to plot PDF.
        plot_cdf: Whether to plot CDF.
        **kwargs: The parameters for the distribution.

    Returns:
        A Plotly figure or an error message.
    """
    try:
        # Validate and parse parameters
        params, error = validate_and_parse_params(distribution_name, kwargs)
        if error:
            return error

        # Create distribution
        distribution = create_distribution(distribution_name, params)

        # Generate plot
        fig, plot_error = plot_distribution(distribution, plot_pdf, plot_cdf)
        if plot_error:
            return plot_error

        return fig

    except Exception as e:
        return f"An error occurred: {e}"


def main() -> None:
    """
    Sets up and launches the Gradio interface.
    """
    with gr.Blocks() as demo:
        gr.Markdown("# PyTorch Distributions Explorer")
        gr.Markdown(
            "Select a distribution, provide its parameters, and visualize its PDF and CDF."
        )

        with gr.Row():
            distribution = gr.Dropdown(
                choices=list(DISTRIBUTIONS.keys()),
                value="Normal",
                label="Distribution",
            )
            plot_pdf = gr.Checkbox(
                label="Plot PDF", value=True
            )
            plot_cdf = gr.Checkbox(
                label="Plot CDF", value=True
            )

        with gr.Row():
            parameter_inputs = gr.Column()

        with gr.Row():
            plot_output = gr.Plot(label="Distribution Plot")

        def on_distribution_change(name: str):
            return update_parameters(name)

        distribution.change(
            on_distribution_change,
            inputs=distribution,
            outputs=parameter_inputs,
        )

        # Initialize parameter inputs
        parameter_components = get_parameter_components(distribution.value)
        parameter_inputs.render = lambda: parameter_components
        for component in parameter_components:
            parameter_inputs.add(component)

        # Define the submit button
        submit = gr.Button("Generate Plot")

        # Define inputs for the plot_function
        def gather_inputs(*args):
            params = {}
            for arg in args:
                # The order of args corresponds to parameter_components
                if isinstance(arg, list):
                    params[arg[0]] = arg[1]
                else:
                    pass
            return params

        submit.click(
            fn=plot_function,
            inputs=[distribution, plot_pdf, plot_cdf] + parameter_components,
            outputs=plot_output,
        )

    demo.launch()


if __name__ == "__main__":
    main()