#!/usr/bin/env python
# -*- encoding: utf-8 -*-

"""
@File : category_feat.py
@Author : Braylon1002
@Version : 1.0
@Contact : S.Braylon1002@gmail.com
Desc : None
"""
from base import BaseFeature

class Category(BaseFeature):
    """
    class for category feature
    """
    def __init__(self, name, column_flow,
                 embedding_name=None, embedding_size=None,
                 vocab_size=None):
        super().__init__(name, column_flow)
        self.embedding_name = embedding_name if embedding_name else name
        self.embedding_size = embedding_size
        self.vocab_size = vocab_size

    def dimension(self):
        if self.column_flow:
            return self.column_flow.transformers[-1].dimension()
        else:
            if self.vocab_size:
                return self.vocab_size
            else:
                raise RuntimeError(
                    "missing dimension which depends on column_flow or vocab_size"
                )