from dataclasses import dataclass
from . import base as base
from jiant.tasks.lib.templates.shared import labels_to_bimap


@dataclass
class Example(base.Example):
    @property
    def label_to_id(self):
        return SentEvalWordContentTask.LABEL_TO_ID


@dataclass
class TokenizedExample(base.TokenizedExample):
    pass


@dataclass
class DataRow(base.DataRow):
    pass


@dataclass
class Batch(base.Batch):
    pass


class SentEvalWordContentTask(base.BaseSentEvalTask):
    Example = Example
    TokenizedExample = TokenizedExample
    DataRow = DataRow
    Batch = Batch

    # SentEval predicts the presence of words in a sentence with 1k-way classification.
    # All examples satisfy:
    #     word in text.lower()
    LABELS = [
        "abandoned",
        "abruptly",
        "accent",
        "access",
        "according",
        "account",
        "ache",
        "ached",
        "acted",
        "acting",
        "actions",
        "actual",
        "address",
        "advantage",
        "advice",
        "afford",
        "agent",
        "agreement",
        "aiden",
        "airport",
        "alarm",
        "albert",
        "alert",
        "alexander",
        "alice",
        "alley",
        "allowing",
        "aloud",
        "amanda",
        "amber",
        "american",
        "amused",
        "amusement",
        "ancient",
        "andy",
        "annie",
        "announced",
        "annoyed",
        "answering",
        "anticipation",
        "anxious",
        "anybody",
        "apologize",
        "appearance",
        "appears",
        "approach",
        "approaching",
        "arched",
        "argue",
        "argument",
        "aria",
        "arrive",
        "ashe",
        "assume",
        "assured",
        "attached",
        "attacked",
        "attempted",
        "audience",
        "available",
        "awful",
        "awkward",
        "background",
        "backpack",
        "backward",
        "backwards",
        "bags",
        "balance",
        "bank",
        "barn",
        "bars",
        "bastard",
        "bath",
        "beating",
        "begged",
        "begins",
        "begun",
        "behavior",
        "bell",
        "belly",
        "belong",
        "belonged",
        "bench",
        "biggest",
        "bike",
        "billy",
        "bird",
        "birds",
        "birth",
        "birthday",
        "bitch",
        "bitter",
        "blank",
        "blast",
        "bleeding",
        "blind",
        "blocked",
        "blond",
        "blonde",
        "bond",
        "bone",
        "bored",
        "bothered",
        "bound",
        "bowed",
        "bowl",
        "boxes",
        "branches",
        "brave",
        "bread",
        "breast",
        "breaths",
        "brian",
        "brick",
        "brilliant",
        "broad",
        "brows",
        "brush",
        "brushing",
        "buddy",
        "build",
        "buildings",
        "bullet",
        "bunch",
        "butt",
        "cage",
        "cain",
        "caleb",
        "callum",
        "calmly",
        "cameron",
        "capable",
        "cards",
        "career",
        "carol",
        "caroline",
        "carriage",
        "cash",
        "cassidy",
        "casual",
        "catching",
        "causing",
        "centuries",
        "chain",
        "chairs",
        "challenge",
        "chamber",
        "chances",
        "changes",
        "changing",
        "charged",
        "checking",
        "chicken",
        "chief",
        "chill",
        "chloe",
        "chocolate",
        "choked",
        "chosen",
        "christian",
        "chuckle",
        "cigarette",
        "circles",
        "circumstances",
        "claim",
        "claimed",
        "clary",
        "classes",
        "claws",
        "clay",
        "cleaned",
        "cleaning",
        "clearing",
        "clicked",
        "cliff",
        "climb",
        "climbing",
        "closely",
        "closest",
        "cloth",
        "clothing",
        "cloud",
        "clue",
        "clung",
        "clutched",
        "coach",
        "cocked",
        "code",
        "colin",
        "collapsed",
        "collar",
        "colors",
        "commander",
        "comment",
        "complicated",
        "concrete",
        "condition",
        "confidence",
        "confident",
        "confirmed",
        "connected",
        "connor",
        "convince",
        "convinced",
        "cook",
        "cops",
        "copy",
        "corners",
        "correct",
        "cost",
        "courage",
        "cousin",
        "covering",
        "covers",
        "crack",
        "cracked",
        "crap",
        "crash",
        "crashed",
        "crawled",
        "cream",
        "create",
        "credit",
        "crept",
        "crime",
        "crossing",
        "crowded",
        "cruel",
        "crystal",
        "cupped",
        "curiosity",
        "current",
        "curse",
        "cursed",
        "cute",
        "cutting",
        "daemon",
        "dagger",
        "damned",
        "damp",
        "danced",
        "danny",
        "dante",
        "darkened",
        "darling",
        "darted",
        "dating",
        "dave",
        "dawn",
        "dealing",
        "delicate",
        "delicious",
        "della",
        "demons",
        "deny",
        "department",
        "desert",
        "deserved",
        "destroy",
        "destroyed",
        "detail",
        "detective",
        "devil",
        "devon",
        "digging",
        "dining",
        "direct",
        "directions",
        "disappear",
        "disappointed",
        "disappointment",
        "disbelief",
        "discuss",
        "disgust",
        "display",
        "distant",
        "distracted",
        "dogs",
        "dollars",
        "double",
        "doubted",
        "dozen",
        "drag",
        "dragging",
        "drake",
        "drank",
        "drawer",
        "drinks",
        "driveway",
        "dropping",
        "drops",
        "drug",
        "drugs",
        "ducked",
        "duke",
        "dull",
        "duncan",
        "duty",
        "dylan",
        "eager",
        "eased",
        "east",
        "eaten",
        "echo",
        "echoed",
        "eddie",
        "edges",
        "elbow",
        "elena",
        "ellie",
        "embarrassed",
        "embrace",
        "emerged",
        "emergency",
        "emotional",
        "ends",
        "enemies",
        "enormous",
        "envelope",
        "equipment",
        "evan",
        "event",
        "everybody",
        "exact",
        "exchange",
        "exchanged",
        "exclaimed",
        "exhausted",
        "exist",
        "existed",
        "existence",
        "exit",
        "expensive",
        "experienced",
        "explanation",
        "exploded",
        "exposed",
        "extended",
        "extremely",
        "eyed",
        "fabric",
        "fade",
        "fairly",
        "faith",
        "fake",
        "falls",
        "families",
        "fangs",
        "farm",
        "farther",
        "favor",
        "feared",
        "feed",
        "fellow",
        "fence",
        "fierce",
        "fifty",
        "file",
        "filling",
        "fingertips",
        "finn",
        "fired",
        "fixed",
        "flame",
        "flashlight",
        "flicked",
        "floating",
        "flushed",
        "foolish",
        "football",
        "footsteps",
        "forcing",
        "formed",
        "frame",
        "freedom",
        "french",
        "friday",
        "friendly",
        "frightened",
        "frown",
        "frowning",
        "froze",
        "frozen",
        "frustrated",
        "fuck",
        "fucking",
        "furious",
        "furniture",
        "fury",
        "gabe",
        "games",
        "garage",
        "garrett",
        "gary",
        "gasp",
        "gates",
        "gather",
        "gathering",
        "gavin",
        "gear",
        "gesture",
        "gestured",
        "ghost",
        "girlfriend",
        "glances",
        "glare",
        "glowing",
        "goodbye",
        "gorgeous",
        "government",
        "gown",
        "grabbing",
        "graham",
        "grand",
        "grandfather",
        "grandma",
        "grant",
        "grasp",
        "grave",
        "greeted",
        "grief",
        "grinning",
        "groan",
        "growl",
        "guessed",
        "guest",
        "guests",
        "guns",
        "gwen",
        "halfway",
        "hank",
        "hanna",
        "happily",
        "happiness",
        "harry",
        "harsh",
        "heal",
        "heartbeat",
        "heaven",
        "height",
        "hero",
        "hesitation",
        "hissed",
        "hitting",
        "holds",
        "holiday",
        "holly",
        "holy",
        "honor",
        "hopefully",
        "horizon",
        "houses",
        "humor",
        "hundreds",
        "hunger",
        "hunt",
        "hunting",
        "hurting",
        "ideas",
        "idiot",
        "ignoring",
        "images",
        "imagination",
        "impressed",
        "impression",
        "inches",
        "including",
        "incredible",
        "incredibly",
        "informed",
        "injured",
        "inner",
        "insane",
        "intended",
        "intensity",
        "invisible",
        "invited",
        "iron",
        "isaac",
        "issue",
        "items",
        "jamie",
        "jared",
        "jealous",
        "jeremy",
        "jerk",
        "jess",
        "jesus",
        "jimmy",
        "joey",
        "johnny",
        "jonas",
        "jordan",
        "joseph",
        "joshua",
        "josie",
        "journey",
        "judge",
        "jumping",
        "kane",
        "karen",
        "katherine",
        "katie",
        "keeps",
        "kevin",
        "kicking",
        "kidding",
        "killer",
        "kisses",
        "knelt",
        "knocking",
        "ladies",
        "landing",
        "language",
        "larger",
        "lately",
        "laura",
        "lauren",
        "lawyer",
        "lean",
        "leapt",
        "learning",
        "lesson",
        "letters",
        "licked",
        "lies",
        "lifting",
        "lightning",
        "likes",
        "lined",
        "liquid",
        "lobby",
        "location",
        "london",
        "lonely",
        "louder",
        "lover",
        "loves",
        "loving",
        "lucien",
        "lucy",
        "major",
        "mama",
        "manner",
        "marks",
        "martin",
        "mask",
        "mason",
        "mass",
        "massive",
        "mate",
        "material",
        "mattered",
        "matters",
        "meaning",
        "meat",
        "medical",
        "melissa",
        "member",
        "members",
        "mental",
        "mere",
        "merely",
        "midnight",
        "military",
        "minds",
        "miranda",
        "mist",
        "mixed",
        "moaned",
        "monday",
        "moonlight",
        "morgan",
        "mortal",
        "motioned",
        "mountains",
        "movements",
        "moves",
        "movies",
        "mumbled",
        "muscle",
        "mystery",
        "nails",
        "named",
        "nathan",
        "nearest",
        "necklace",
        "needing",
        "nerves",
        "nervously",
        "nicolas",
        "nicole",
        "nightmare",
        "nights",
        "nina",
        "nodding",
        "nods",
        "nora",
        "noted",
        "notes",
        "numbers",
        "object",
        "occurred",
        "offering",
        "officers",
        "oliver",
        "opens",
        "opinion",
        "option",
        "orange",
        "ordinary",
        "original",
        "ourselves",
        "overhead",
        "owner",
        "packed",
        "pages",
        "painful",
        "paint",
        "painted",
        "painting",
        "palace",
        "palms",
        "paris",
        "parted",
        "particularly",
        "partner",
        "parts",
        "passenger",
        "passion",
        "patch",
        "patience",
        "patient",
        "patrick",
        "patted",
        "paying",
        "period",
        "permission",
        "pete",
        "photo",
        "pile",
        "pillow",
        "pissed",
        "pity",
        "placing",
        "plain",
        "planet",
        "plastic",
        "pleasant",
        "plus",
        "pockets",
        "points",
        "polite",
        "popped",
        "positive",
        "possibility",
        "potential",
        "pounded",
        "powers",
        "precious",
        "prefer",
        "prepare",
        "president",
        "preston",
        "pretend",
        "pretending",
        "previous",
        "price",
        "pride",
        "princess",
        "prison",
        "professor",
        "program",
        "progress",
        "project",
        "proof",
        "proper",
        "property",
        "protection",
        "protest",
        "provide",
        "provided",
        "pulse",
        "punch",
        "punched",
        "pure",
        "purple",
        "puts",
        "quinn",
        "quit",
        "racing",
        "radio",
        "rafe",
        "raise",
        "raising",
        "rapidly",
        "rare",
        "rarely",
        "reaches",
        "react",
        "realizing",
        "rear",
        "rebecca",
        "recall",
        "received",
        "recently",
        "record",
        "reflection",
        "refuse",
        "regular",
        "reluctantly",
        "remaining",
        "remains",
        "remembering",
        "remind",
        "remove",
        "replaced",
        "request",
        "required",
        "rescue",
        "research",
        "resist",
        "responsibility",
        "responsible",
        "result",
        "returning",
        "reveal",
        "revealed",
        "revealing",
        "ribs",
        "riding",
        "rifle",
        "roar",
        "robe",
        "robert",
        "rode",
        "rope",
        "rubbing",
        "ruin",
        "ruined",
        "rule",
        "runs",
        "rushing",
        "sadly",
        "sadness",
        "sake",
        "sally",
        "satisfaction",
        "satisfied",
        "saturday",
        "savannah",
        "saving",
        "scanned",
        "scar",
        "scare",
        "scattered",
        "scott",
        "scrambled",
        "screams",
        "seated",
        "seats",
        "secrets",
        "section",
        "sees",
        "self",
        "sell",
        "sensation",
        "sensed",
        "senses",
        "sentence",
        "separate",
        "serena",
        "series",
        "serve",
        "served",
        "settle",
        "sexy",
        "shade",
        "shakes",
        "shame",
        "sharply",
        "shattered",
        "sheet",
        "sheets",
        "sheriff",
        "shield",
        "shining",
        "shiver",
        "shivered",
        "shooting",
        "shopping",
        "shore",
        "shortly",
        "shorts",
        "shout",
        "shouting",
        "shown",
        "shrug",
        "shuddered",
        "sidewalk",
        "sideways",
        "signal",
        "signs",
        "silk",
        "silly",
        "similar",
        "sisters",
        "site",
        "sits",
        "skills",
        "skirt",
        "skull",
        "slapped",
        "sleeve",
        "slide",
        "slipping",
        "slowed",
        "slumped",
        "smaller",
        "smith",
        "snap",
        "social",
        "society",
        "sofa",
        "somewhat",
        "sooner",
        "sophia",
        "sophie",
        "souls",
        "sounding",
        "source",
        "spare",
        "speech",
        "spencer",
        "spending",
        "spinning",
        "split",
        "spring",
        "squeeze",
        "stands",
        "stated",
        "statement",
        "states",
        "steal",
        "stepping",
        "stiff",
        "stirred",
        "stole",
        "stones",
        "stops",
        "straightened",
        "stream",
        "stretch",
        "strike",
        "strode",
        "stroked",
        "struggle",
        "struggling",
        "student",
        "students",
        "studying",
        "stuffed",
        "style",
        "success",
        "suggest",
        "sunday",
        "sunlight",
        "supplies",
        "surprisingly",
        "surrounding",
        "survived",
        "suspect",
        "suspected",
        "suspicious",
        "swallow",
        "sweetheart",
        "swiftly",
        "swing",
        "switch",
        "t-shirt",
        "tables",
        "tail",
        "tapped",
        "tara",
        "target",
        "task",
        "tasted",
        "taught",
        "taylor",
        "teach",
        "teacher",
        "teased",
        "teasing",
        "television",
        "tells",
        "temper",
        "temple",
        "tense",
        "tent",
        "terms",
        "terrified",
        "terror",
        "tessa",
        "text",
        "thigh",
        "thighs",
        "thousands",
        "threatened",
        "threatening",
        "throughout",
        "thrown",
        "thrust",
        "tighter",
        "till",
        "tipped",
        "toby",
        "toes",
        "tony",
        "tore",
        "torn",
        "total",
        "tough",
        "towel",
        "tower",
        "trace",
        "tracks",
        "traffic",
        "trailed",
        "trained",
        "trap",
        "trapped",
        "traveled",
        "travis",
        "tray",
        "treat",
        "treated",
        "trembled",
        "trevor",
        "trick",
        "tristan",
        "trunk",
        "trusted",
        "twelve",
        "ugly",
        "unconscious",
        "unexpected",
        "uniform",
        "universe",
        "unknown",
        "unlike",
        "unsure",
        "unusual",
        "upper",
        "upright",
        "upward",
        "useful",
        "useless",
        "valley",
        "vanished",
        "various",
        "vehicle",
        "veins",
        "victor",
        "video",
        "vincent",
        "violet",
        "visible",
        "vulnerable",
        "walks",
        "wandered",
        "warrior",
        "wash",
        "washed",
        "waving",
        "weather",
        "werewolf",
        "whipped",
        "whispers",
        "whom",
        "willow",
        "winced",
        "wings",
        "winked",
        "wise",
        "wishing",
        "witch",
        "wolves",
        "worn",
        "wounded",
        "wounds",
        "wrap",
        "wrapping",
        "wrists",
        "wrote",
        "xavier",
        "yanked",
        "yard",
        "yards",
        "yelling",
        "zach",
    ]
    LABEL_TO_ID, ID_TO_LABEL = labels_to_bimap(LABELS)
