from typing import Union

import torch.nn as nn

from .actor import Actor
from .critic import Critic
from .reward_model import RewardModel


def get_base_model(model: Union[Actor, Critic, RewardModel]) -> nn.Module:
    """Get the base model of our wrapper classes.
    For Actor, Critic and RewardModel, return ``model.model``, 
    it's usually a ``transformers.PreTrainedModel``.

    Args:
        model (nn.Module): model to get base model from

    Returns:
        nn.Module: the base model
    """
    assert isinstance(model, (Actor, Critic, RewardModel)), \
        f'Expect Actor, Critic or RewardModel, got {type(model)}, use unwrap_model first.'
    return model.model


__all__ = ['Actor', 'Critic', 'RewardModel', 'get_base_model']