"""
This type stub file was generated by pyright.
"""

from .dgl_dataset import DGLBuiltinDataset

"""Tree-structured data.
Including:
    - Stanford Sentiment Treebank
"""
__all__ = ['SST', 'SSTDataset']
class SSTDataset(DGLBuiltinDataset):
    r"""Stanford Sentiment Treebank dataset.

    .. deprecated:: 0.5.0
        
        - ``trees`` is deprecated, it is replaced by:

            >>> dataset = SSTDataset()
            >>> for tree in dataset:
            ....    # your code here

        - ``num_vocabs`` is deprecated, it is replaced by ``vocab_size``.

    Each sample is the constituency tree of a sentence. The leaf nodes
    represent words. The word is a int value stored in the ``x`` feature field.
    The non-leaf node has a special value ``PAD_WORD`` in the ``x`` field.
    Each node also has a sentiment annotation: 5 classes (very negative,
    negative, neutral, positive and very positive). The sentiment label is a
    int value stored in the ``y`` feature field.
    Official site: `<http://nlp.stanford.edu/sentiment/index.html>`_

    Statistics:

    - Train examples: 8,544
    - Dev examples: 1,101
    - Test examples: 2,210
    - Number of classes for each node: 5

    Parameters
    ----------
    mode : str, optional
        Should be one of ['train', 'dev', 'test', 'tiny']
        Default: train
    glove_embed_file : str, optional
        The path to pretrained glove embedding file.
        Default: None
    vocab_file : str, optional
        Optional vocabulary file. If not given, the default vacabulary file is used.
        Default: None
    raw_dir : str
        Raw file directory to download/contains the input data directory.
        Default: ~/.dgl/
    force_reload : bool
        Whether to reload the dataset. Default: False
    verbose: bool
        Whether to print out progress information. Default: True.

    Attributes
    ----------
    vocab : OrderedDict
        Vocabulary of the dataset
    trees : list
        A list of DGLGraph objects
    num_classes : int
        Number of classes for each node
    pretrained_emb: Tensor
        Pretrained glove embedding with respect the vocabulary.
    vocab_size : int
        The size of the vocabulary
    num_vocabs : int
        The size of the vocabulary

    Notes
    -----
    All the samples will be loaded and preprocessed in the memory first.

    Examples
    --------
    >>> # get dataset
    >>> train_data = SSTDataset()
    >>> dev_data = SSTDataset(mode='dev')
    >>> test_data = SSTDataset(mode='test')
    >>> tiny_data = SSTDataset(mode='tiny')
    >>>
    >>> len(train_data)
    8544
    >>> train_data.num_classes
    5
    >>> glove_embed = train_data.pretrained_emb
    >>> train_data.vocab_size
    19536
    >>> train_data[0]
    Graph(num_nodes=71, num_edges=70,
      ndata_schemes={'x': Scheme(shape=(), dtype=torch.int64), 'y': Scheme(shape=(), dtype=torch.int64), 'mask': Scheme(shape=(), dtype=torch.int64)}
      edata_schemes={})
    >>> for tree in train_data:
    ...     input_ids = tree.ndata['x']
    ...     labels = tree.ndata['y']
    ...     mask = tree.ndata['mask']
    ...     # your code here
    """
    PAD_WORD = ...
    UNK_WORD = ...
    def __init__(self, mode=..., glove_embed_file=..., vocab_file=..., raw_dir=..., force_reload=..., verbose=...) -> None:
        ...
    
    def process(self): # -> None:
        ...
    
    def has_cache(self): # -> bool:
        ...
    
    def save(self): # -> None:
        ...
    
    def load(self): # -> None:
        ...
    
    @property
    def trees(self): # -> list[Unknown]:
        ...
    
    @property
    def vocab(self): # -> OrderedDict[Unknown, Unknown] | Any:
        r""" Vocabulary

        Returns
        -------
        OrderedDict
        """
        ...
    
    @property
    def pretrained_emb(self): # -> Any | None:
        r"""Pre-trained word embedding, if given."""
        ...
    
    def __getitem__(self, idx):
        r""" Get graph by index

        Parameters
        ----------
        idx : int

        Returns
        -------
        :class:`dgl.DGLGraph`

            graph structure, word id for each node, node labels and masks.

            - ``ndata['x']``: word id of the node
            - ``ndata['y']:`` label of the node
            - ``ndata['mask']``: 1 if the node is a leaf, otherwise 0
        """
        ...
    
    def __len__(self): # -> int:
        r"""Number of graphs in the dataset."""
        ...
    
    @property
    def num_vocabs(self): # -> int:
        ...
    
    @property
    def vocab_size(self): # -> int:
        r"""Vocabulary size."""
        ...
    
    @property
    def num_classes(self): # -> Literal[5]:
        r"""Number of classes for each node."""
        ...
    


SST = SSTDataset
