# @Time     : 2024/7/19 14:27
# @Author   : Hang Zhou
# @Email    : fjklqq@gmail.com
# @Software : Python 3.11
# @About    :
__all__ = ["TokenClassificationTaskArguments", "TokenClassificationPredictArguments"]

from dataclasses import dataclass, field
from typing import Optional

from tt4l.factory.base.arguments import BaseTaskArguments, BasePredictArguments


@dataclass
class TokenClassificationTaskArguments(BaseTaskArguments):
    token_column_name: Optional[str] = field(
        default='token', metadata={"help": "The column name of tokens to input in the file (a csv or JSON file)."}
    )
    label_column_name: Optional[str] = field(
        default='label', metadata={"help": "The column name of label to input in the file (a csv or JSON file)."}
    )
    task_name: Optional[str] = field(default="ner", metadata={"help": "The name of the task (ner, pos...)."})
    label_all_tokens: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to put the label for one word on all tokens of generated by that word or just on the "
                "one (in which case the other tokens will have a padding index)."
            )
        },
    )
    return_entity_level_metrics: bool = field(
        default=False,
        metadata={"help": "Whether to return all the entity levels during evaluation or just the overall ones."},
    )


@dataclass
class TokenClassificationPredictArguments(BasePredictArguments):
    text_column_name: Optional[str] = field(
        default='sentence', metadata={"help": "The column name of text to input in the file (a csv or JSON file)."}
    )
