from typing import Optional

TEXT2FOLD = {
"REG_SOCIAL_ANIME_PLANET_RATING": 0,
"REG_CONSUMER_JC_PENNEY_PRODUCT_PRICE": 0,
"MUL_CONSUMER_PRODUCT_SENTIMENT": 0,
"MUL_SOCIAL_NEWS_CHANNEL_CATEGORY": 0,
"REG_HOUSES_CALIFORNIA_PRICES_2020": 0,
"REG_CONSUMER_BABIES_R_US_PRICES": 0,
"BIN_SOCIAL_JIGSAW_TOXICITY": 0,
"REG_SOCIAL_MUSEUMS_US_REVENUES": 0,
"REG_PROFESSIONAL_EMPLOYEE_RENUMERATION_VANCOUBER": 0,
"REG_SOCIAL_SPOTIFY_POPULARITY": 0,
"MUL_SOCIAL_SPOTIFY_GENRES": 0,             # For TabSTAR-v2's Text evaluation
"REG_FOOD_RAMEN_RATINGS_2022": 1,
"REG_SOCIAL_FILMTV_MOVIE_RATING_ITALY": 1,
"REG_CONSUMER_AMERICAN_EAGLE_PRICES": 1,
"REG_PROFESSIONAL_EMPLOYEE_SALARY_MONTGOMERY": 1,
"MUL_CONSUMER_WOMEN_ECOMMERCE_CLOTHING_REVIEW": 1,
"REG_SPORTS_FIFA22_WAGES": 1,
"REG_FOOD_WINE_VIVINO_SPAIN": 1,
"MUL_FOOD_WINE_REVIEW": 1,
"REG_PROFESSIONAL_COMPANY_EMPLOYEES_SIZE": 1,
"MUL_FOOD_YELP_REVIEWS": 1,
"REG_CONSUMER_BIKE_PRICE_BIKEWALE": 2,
"REG_TRANSPORTATION_USED_CAR_PAKISTAN": 2,
"MUL_HOUSES_MELBOURNE_AIRBNB": 2,
"BIN_PROFESSIONAL_KICKSTARTER_FUNDING": 2,
"MUL_FOOD_MICHELIN_GUIDE_RESTAURANTS": 2,
"REG_FOOD_ALCOHOL_WIKILIQ_PRICES": 2,
"REG_FOOD_COFFEE_REVIEW": 2,
"BIN_SOCIAL_IMDB_GENRE_PREDICTION": 2,
"REG_SOCIAL_KOREAN_DRAMA": 2,
"REG_FOOD_BEER_RATINGS": 2,
"REG_SOCIAL_BOOK_READABILITY_CLEAR": 3,
"MUL_SOCIAL_GOOGLE_QA_TYPE_REASON": 3,
"REG_SOCIAL_VIDEO_GAMES_SALES": 3,
"REG_SOCIAL_MOVIES_DATASET_REVENUE": 3,
"REG_FOOD_WINE_POLISH_MARKET_PRICES": 3,
"REG_FOOD_CHOCOLATE_BAR_RATINGS": 3,
"REG_TRANSPORTATION_USED_CAR_MERCEDES_BENZ_ITALY": 3,
"REG_CONSUMER_BOOK_PRICE_PREDICTION": 3,
"REG_TRANSPORTATION_USED_CAR_SAUDI_ARABIA": 3,
"MUL_TRANSPORTATION_US_ACCIDENTS_MARCH23": 3,
"REG_FOOD_ZOMATO_RESTAURANTS": 4,
"REG_SOCIAL_MOVIES_ROTTEN_TOMATOES": 4,
"REG_CONSUMER_CAR_PRICE_CARDEKHO": 4,
"REG_SPORTS_NBA_DRAFT_VALUE_OVER_REPLACEMENT": 4,
"MUL_PROFESSIONAL_DATA_SCIENTIST_SALARY": 4,
"REG_SOCIAL_BOOKS_GOODREADS": 4,
"REG_PROFESSIONAL_ML_DS_AI_JOBS_SALARIES": 4,
"REG_CONSUMER_MERCARI_ONLINE_MARKETPLACE": 4,
"BIN_PROFESSIONAL_FAKE_JOB_POSTING": 4,
"REG_PROFESSIONAL_SCIMAGOJR_ACADEMIC_IMPACT": 4,
}


PRETRAIN2FOLD = {
"MUL_ANONYM_FABERT": 0,
"BIN_CONSUMER_DRESS_DALES": 0,
"REG_SOCIAL_OCCUPATION_MOBILITY_SOCMOB": 0,
"BIN_HEALTHCARE_LIVER_INDIAN_ILPD": 0,
"BIN_NATURE_OZONE_LEVEL": 0,
"MUL_SOCIAL_BIAS_FRAMES": 0,
"BIN_COMPUTERS_PC4_CODE_DEFECTIONS": 0,
"REG_SPORTS_FOOTBALL_MANAGER_IMPORTANT_MATCHES": 0,
"MUL_SOCIAL_OKCUPID_DATING_JOB_STEM": 0,
"MUL_COMPUTERS_IMAGE_OPT_DIGITS": 0,
"BIN_ANONYM_MADELONE": 0,
"MUL_HEALTHCARE_HYPOTHYROID": 0,
"MUL_SPORTS_ACTIVITY_LOCALIZED_LDPA": 0,
"MUL_SCIENCE_ANNEAL_CHEMICAL": 0,
"REG_COMPUTERS_META_LEVEL_LEARNING": 0,
"REG_HOUSES_PERTH": 0,
"MUL_CONSUMER_ELECTRICITY_TAMILNADU": 0,
"REG_HOUSES_SEATTLE_KINGS_COUNTY": 0,
"REG_HOUSES_BOSTON_PRICE_NOMINAL": 0,
"MUL_SCIENCE_VOWEL_SPEAKER_RECOGNITION": 0,
"BIN_CONSUMER_MOBILE_CHURN": 0,
"REG_TRANSPORTATION_US_BIKE_SHARING_DEMAND": 0,
"BIN_ANONYM_JASMINE": 0,
"BIN_GEOGRAPHY_NOMAO_SEARCH_ENGINE": 0,
"REG_COMPUTERS_ROBOT_SARCOS": 0,
"REG_SCIENCE_SIMULATION_2D_PLANES": 0,
"BIN_SPORTS_RUN_OR_WALK": 0,
"REG_ANONYM_BUZZ_IN_SOCIAL_MEDIA_TWITTER": 0,
"BIN_SOCIAL_COMPASS_TWO_YEARS_OFFEND": 0,
"MUL_FINANCIAL_P2P_LOAN": 0,
"MUL_FINANCIAL_CREDIT_RISK_CHINA": 0,
"REG_NATURE_WIND_SPEED_IRELAND": 0,
"MUL_ANONYM_SYNTHETIC_CONTROL": 0,
"REG_HOUSES_GERMAN_HOUSES": 0,
"BIN_FINANCIAL_BANK_ACCOUNT_FRAUD_BAF": 0,
"MUL_GENETICS_ACP_BREAST_CANCER": 0,
"REG_ANONYM_PLANTS_TEXTURE": 0,
"REG_NATURE_FOREST_FIRES": 0,
"MUL_ANONYM_PENDIGITS": 0,
"BIN_FINANCIAL_ACCOUNT_DEBT_SOUTH_AFRICA": 0,
"BIN_COMPUTERS_KC1_CODE_DEFECTIONS": 0,
"BIN_COMPUTERS_JM1_CODE_DEFECTIONS": 0,         # For TabSTAR-v2's TabArena evaluation
"BIN_TRANSPORTATION_ROAD_SAFETY_GENDER": 0,
"BIN_ANONYM_GINA": 0,
"REG_ANONYM_YPROP": 0,
"MUL_TRANSPORTATION_CAR_STUDENT_INSURANCE_RISK": 0,
"BIN_FINANCIAL_BANK_MARKETING": 0,
"MUL_TRANSPORTATION_VEHICLE_SILHOUETTE": 0,
"REG_HOUSES_BOSTON_HOUSE": 0,
"BIN_CONSUMER_CLICK_PREDICTION": 0,
"MUL_SCIENCE_DRUG_DIRECTORY_FDA_UNAPPROVED": 0,
"MUL_NATURE_FOREST_COVERTYPE": 0,
"REG_COMPUTERS_CPU_ACTIVITY": 0,
"BIN_HEALTHCARE_GLIOMA_BRAIN_TUMOR": 0,
"BIN_ANONYM_AUTOUNIV_AU1": 0,
"BIN_COMPUTERS_PHISHING_URL_WEBSITES": 0,
"BIN_HEALTHCARE_THYROID_CANCER": 0,
"BIN_PROFESSIONAL_LICD_LABOR_RIGHTS": 0,
"MUL_TRANSPORTATION_CAR_ORIGIN_COUNTRY": 0,
"BIN_SOCIAL_SPEED_DATING": 0,
"MUL_ANONYM_ZERNIKE_MFEAT": 0,
"BIN_HEALTHCARE_HEART_DISEASE": 0,
"BIN_ANONYM_MAMMOGRAPHY": 0,
"MUL_HEALTHCARE_CANCER_CERVICAL": 0,
"REG_SCIENCE_SULFUR": 0,
"BIN_CONSUMER_TOUR_TRAVEL_CHURN": 0,
"MUL_SOCIAL_GOT_SCRIPTS": 0,
"REG_COMPUTERS_ROBOT_KIN8NM": 0,
"BIN_SOCIAL_TIC_TAC_TOE": 0,
"MUL_INDUSTRIAL_STEEL_PLATES_FAULT": 0,
"BIN_SCIENCE_PARTICLE_HIGGS": 0,
"BIN_ANONYM_KDDCUP_09_UPSELLING": 0,
"BIN_HEALTHCARE_SEPSIS_PHYSIONET": 0,
"BIN_SOCIAL_TWITTER_DISASTER": 0,
"MUL_NATURE_SOYBEAN_DISEASE": 0,
"BIN_ANONYM_PHILIPPINE": 0,
"REG_PROFESSIONAL_STUDENT_PERFORMANCE_PORTUGAL": 0,
"MUL_PROFESSIONAL_TEACHER_ASSISTANT_EVALUATION_TAE": 1,
"MUL_SPORTS_BASEBALL_HALL_OF_FAME": 1,
"MUL_COMPUTERS_IMAGE_CIFAR10": 1,
"MUL_ANONYM_SPOKEN_ARABIC_DIGIT": 1,
"REG_COMPUTERS_YOUTUBE_VIDEO_TRANSCODING": 1,
"MUL_SOCIAL_WIKIPEDIA_TALK_LABELS_ATTACKS": 1,
"BIN_ANONYM_XD6": 1,
"REG_NATURE_MYANMAR_AIR_QUALITY": 1,
"MUL_TRANSPORTATION_TRAFFIC_ACCIDENTS_FARS": 1,
"MUL_ANONYM_HELENA": 1,
"MUL_HEALTHCARE_EGYPTIAN_SKULLS": 1,
"BIN_NATURE_COLIC_HORSES_SURGICAL_LESION": 1,
"MUL_PROFESSIONAL_NURSERY_APPLICATIONS_SLOVENIA": 1,
"REG_SCIENCE_GRID_STABILITY": 1,
"BIN_HEALTHCARE_MENTAL_HEALTH_TECH": 1,
"BIN_COMPUTERS_MOZILLA": 1,
"BIN_HEALTHCARE_CELLS_WDBC_WISCONSIN_BREAST_CANCER": 1,
"REG_SPORTS_NBA_2K20_PLAYERS_RATING": 1,
"BIN_FINANCIAL_CC_TAIWAN_CREDIT_DEFAULT": 1,
"BIN_NATURE_DISEASED_TREES_WILT": 1,
"BIN_HEALTHCARE_DIABETES_EARLY_STAGE": 1,
"MUL_COMPUTERS_IMAGE_SEGMENTATION": 1,
"MUL_HEALTHCARE_HEPATITIS_C_EGYPT": 1,
"MUL_COMPUTERS_IMAGE_GESTURE_PHASE_SEGMENTATION": 1,
"MUL_NATURE_AIR_QUALITY_POLLUTION": 1,
"REG_ANONYM_TOPO": 1,
"BIN_GENETICS_MUSK_MOLECULES": 1,
"BIN_ANONYM_ARCENE": 1,
"REG_HOUSES_IOWA_HOUSES_PRICES": 1,
"BIN_CONSUMER_ELECTRICITY_PRICE_TREND": 1,
"BIN_SPORTS_PRO_FOOTBALL_BETS_PROFB": 1,
"BIN_CONSUMER_ONLINE_SHOPPERS_PURCHASE_INTENTION": 1,
"MUL_ANONYM_FIRST_ORDER_THEOREM_PROVING": 1,
"BIN_COMPUTERS_PHISHING_WEBSITE_HUDDERSFIELD": 1,
"MUL_COMPUTERS_PHISHING_WEBSITE_HUDDERSFIELD": 1,         # For TabSTAR-v2's TabArena evaluation
"MUL_SCIENCE_MICE_PROTEIN": 1,
"BIN_TRANSPORTATION_ESTONIA_DISASTER_PASSENGERS": 1,
"MUL_NATURE_EUCALYPTUS_SEED": 1,
"BIN_TRANSPORTATION_TITANIC_SURVIVAL": 1,
"MUL_SOCIAL_SOCC_OPINION_COMMENTS_CORPUS_TOXICITY": 1,
"MUL_FINANCIAL_LOAN_IBRD": 1,
"MUL_CONSUMER_INTERNET_USAGE_PROFESSION": 1,
"REG_NATURE_CLIMATE_CHANGE_AGRICULTURE_FINANCIAL": 1,
"BIN_CONSUMER_CHURN_TELCO_EUROPA": 1,
"MUL_FOOD_WINE_QUALITY_CAT": 1,
"REG_FOOD_WINE_QUALITY": 1,                             # For TabSTAR-v2's TabArena evaluation
"BIN_ANONYM_TWONORM": 1,
"BIN_ANONYM_GAMETES": 1,
"REG_HOUSES_MIAMI": 1,
"BIN_SCIENCE_HIV_QSAR": 1,
"MUL_SCIENCE_HIV_QSAR": 1,                              # For TabSTAR-v2's TabArena evaluation
"REG_SOCIAL_STRIKES_PER_COUNTRY": 1,
"MUL_TRANSPORTATION_SHUTTLE_SPACE": 1,
"MUL_ANONYM_DNA": 1,
"REG_TRANSPORTATION_US_AIRPORT_PASSENGERS": 1,
"REG_COMPUTERS_PUMA_ROBOT_ARM": 1,
"MUL_HEALTHCARE_MENTAL_HEALTH_OCCUPATION": 1,
"REG_TRANSPORTATION_FLIGHT_SIMULATION_ELEVATORS": 1,
"BIN_SCIENCE_EYE_STATE_EEG": 1,
"BIN_ANONYM_NUMERAI_28_6": 1,
"REG_SOCIAL_OKCUPID_DATING_PROFILE_AGE": 1,
"REG_NATURE_NO2_POLLUTION_NORWAY": 1,
"REG_SCIENCE_AUCTION_VERIFICATION": 1,
"REG_NATURE_ABALONE_FISH_RINGS": 1,
"MUL_HEALTHCARE_BABY_SLEEP_STATE": 1,
"MUL_SCIENCE_VOWEL_JAPANESE": 1,
"REG_COMPUTERS_ALGO_RUNTIME_SAT11": 1,
"MUL_NATURE_SOLAR_FLARES": 1,
"REG_HOUSES_CALIFORNIA_HOUSES": 1,
"BIN_CONSUMER_WHOLESALE_CUSTOMERS_PORTUGAL": 1,
"REG_HEALTHCARE_DIABETES_SKLEARN": 1,
"MUL_SPORTS_CHESS_JUNGLE": 1,
"MUL_SCIENCE_WAVEFORM_500": 1,
"REG_SCIENCE_QSAR_TID_10980": 1,
"BIN_FINANCIAL_AUDIT_RISK": 1,
"MUL_SOCIAL_HOLISTIC_BIAS": 2,
"BIN_PROFESSIONAL_ACADEMIC_COLLEGE_SCORECARD": 2,
"REG_SCIENCE_AIRFOIL_SELF_NOISE": 2,
"REG_TRANSPORTATION_NAVAL_PROPULSION_PLANT": 2,
"REG_ANONYM_BANK_32NH": 2,
"REG_SCIENCE_WAVE_ENERGY": 2,
"REG_SCIENCE_SUPERCONDUCTIVITY": 2,
"MUL_COMPUTERS_IMAGE_ARTIFICIAL_CHARACTERS": 2,
"MUL_SOCIAL_US_ELECTIONS_SPEECHES": 2,
"MUL_ANONYM_LED24": 2,
"REG_SPORTS_BASEBALL_HITTER_SALARY": 2,
"BIN_NATURE_MUSHROOM_POISONOUS": 2,
"MUL_HEALTHCARE_HEART_CARDIOTOCOGRAPHY": 2,
"REG_HEALTHCARE_BODY_FAT": 2,
"REG_NATURE_EL_NINO_KDD": 2,
"REG_CONSUMER_MEDICAL_CHARGES": 2,
"REG_HOUSES_LISBON": 2,
"MUL_SCIENCE_BALANCE_SCALE": 2,
"MUL_FINANCIAL_CORPORATE_CREDIT_RATING_SYMBOL": 2,
"BIN_ANONYM_KDDCUP_09_APPETENCY": 2,
"REG_SPORTS_FIFA20_PLAYERS_VALUE": 2,
"REG_CONSUMER_DIAMONDS_PRICES": 2,
"BIN_FINANCIAL_BANK_CUSTOMER_CHURN_SHRUTIME": 2,
"BIN_ANONYM_ALBERT": 2,
"REG_FINANCIAL_STOCK_AEROSPACE": 2,
"MUL_HEALTHCARE_PBCSEQ_BILIARY_CIRRHOSIS": 2,
"BIN_FINANCIAL_LT_VEHICLE_LOAN": 2,
"REG_SCIENCE_PHYSIOCHEMICAL_PROTEIN": 2,
"MUL_ANONYM_JANNIS": 2,
"BIN_TRANSPORTATION_AIRLINES_DEPARTURE_DELAY": 2,
"REG_SPORTS_MONEYBALL": 2,
"BIN_HEALTHCARE_BREAST_CANCER_YUGOSLAVIA": 2,
"BIN_SOCIAL_EDUCATIONAL_TRANSITIONS_IRISH": 2,
"MUL_COMPUTERS_INTRUSION_DETECTION_KDD": 2,
"MUL_SOCIAL_BROWN_FROWN_CORPORA_COLLINS": 2,
"MUL_HEALTHCARE_APPENDICITIS_REGENSBURG_PEDIATRIC": 2,
"BIN_FINANCIAL_BANKRUPTCY_QUALITATIVE": 2,
"MUL_ANONYM_DILBERT": 2,
"BIN_HEALTHCARE_HEART_FAILURE": 2,
"BIN_CONSUMER_INTERNET_ADVERTISEMENTS": 2,
"REG_ANONYM_ALLSTATE_CLAIM_SEVERITY": 2,
"REG_SOCIAL_US_CRIME": 2,
"MUL_SOCIAL_HARRY_POTTER_FAN_FICTION": 2,
"REG_SOCIAL_COLLEGES_GRANTS": 2,
"REG_SCIENCE_ENERGY_EFFICIENCY": 2,
"BIN_ANONYM_AUSTRALIAN_CREDIT_APPROVAL": 2,
"BIN_ANONYM_OIL_SPILL": 2,
"MUL_SCIENCE_BATCH_CHORAL_HARMONY": 2,
"BIN_FINANCIAL_BANK_PERSONAL_LOAN_MODELING": 2,
"BIN_FINANCIAL_CC_FRAUD_TRANSACTION": 2,
"BIN_CONSUMER_BLASTCHAR_TELCOM_CHURN": 2,
"REG_FINANCIAL_INSURANCE_HEALTH_HOURS_WORKED_WIFE": 2,
"MUL_COMPUTERS_IMAGE_MNIST_JAPANESE_KUZUSHIJI_49": 2,
"BIN_ANONYM_SYLVINE": 2,
"BIN_FINANCIAL_CREDIT_GERMAN": 2,
"BIN_CONSUMER_HOTEL_REVIEW": 2,
"BIN_SOCIAL_HUMAN_CHOICE_PREDICTION_LM_GAMES": 2,
"BIN_SOCIAL_SPAM_EMAILS_SPAMBASE": 2,
"REG_FINANCIAL_CC_USER_SEGMENTATION": 2,
"MUL_COMPUTERS_PAGE_BLOCK_PARSING": 2,
"MUL_SCIENCE_QSAR_BIOCONCENTRATION": 2,
"MUL_NATURE_CJS_MAPLE_TREES": 2,
"MUL_ANONYM_VOLCANOES_VENUS": 2,
"BIN_FINANCIAL_ADULT_INCOME": 2,
"REG_TRANSPORTATION_NYC_TAXI_TRIP_DURATION": 2,
"BIN_PROFESSIONAL_AMAZON_EMPLOYEE_ACCESS": 2,
"MUL_SOCIAL_FUNPEDIA": 2,
"BIN_PROFESSIONAL_EMPLOYEE_IBM_ATTRITION": 2,
"BIN_ANONYM_PORTO_SEGURO": 2,
"BIN_ANONYM_CHRISTINE": 3,
"REG_SOCIAL_VOTING_SPACE_GEOGRAPHIC_ANALYSIS": 3,
"BIN_ANONYM_ADA": 3,
"BIN_ANONYM_KDDCUP_98_DIRECT_MAIL": 3,
"BIN_HEALTHCARE_BLOOD_TRANSFUSION": 3,
"MUL_SOCIAL_BOOKS_AUTHORSHIP_ANALCATDATA": 3,
"REG_TRANSPORTATION_NYC_TAXI_TIP": 3,
"MUL_TRANSPORTATION_CAR_ACCEPTABILITY": 3,
"REG_NATURE_FISH_TOXICITY": 3,
"BIN_SPORTS_NBA_SHOTS": 3,
"MUL_SCIENCE_MIMIC_ROUTE": 3,
"MUL_GENETICS_SPLICE_DNA": 3,
"MUL_COMPUTERS_IMAGE_MNIST_FASHION": 3,
"BIN_FINANCIAL_TVS_CREDIT_LOAN": 3,
"BIN_ANONYM_GUILLERMO": 3,
"REG_CONSUMER_BOOKING_HOTEL_REVIEW": 3,
"MUL_COMPUTERS_IMAGE_MNIST_DIGITS": 3,
"BIN_TRANSPORTATION_CAR_BAD_BUY_KICK": 3,
"BIN_FINANCIAL_FRAUD_DETECTION": 3,
"REG_ANONYM_DIONIS": 3,
"MUL_TRANSPORTATION_TRAFFIC_VIOLATION": 3,
"BIN_SCIENCE_MAGIC_TELESCOPE": 3,
"BIN_HEALTHCARE_BREAST_CANCER_WISCONSIN": 3,
"MUL_TRANSPORTATION_AIRLINES_TWITTER_SENTIMENT": 3,
"REG_SCIENCE_QSAR_TID_11": 3,
"BIN_HEALTHCARE_DIABETES_US130": 3,     # For TabSTAR-v2's TabArena evaluation
"MUL_HEALTHCARE_DIABETES_US130": 3,
"BIN_SOCIAL_POLICE_INCIDENTS_SAN_FRANCISCO": 3,
"MUL_ANONYM_MICRO_MASS": 3,
"MUL_ANONYM_ROBERT": 3,
"MUL_SOCIAL_SUPREME_COURT_ANALCATDATA": 3,
"BIN_HEALTHCARE_BONE_MARROW_TRANSPLANT_CHILDREN": 3,
"REG_HOUSES_HOUSEHOLD_MONTHLY_ELECTRICITY": 3,
"REG_TRANSPORTATION_CAR_GM_PRICE": 3,
"MUL_NATURE_YEAST_PROTEIN": 3,
"BIN_COMPUTERS_CPU_TOKYO1": 3,
"MUL_HEALTHCARE_AUDIOLOGY_DIAGNOSTIC": 3,
"MUL_COMPUTERS_IMAGE_LETTER_RECOGNITION": 3,
"MUL_SOCIAL_STACKOVERFLOW_POLARITY": 3,
"REG_ANONYM_SANTANDER_TRANSACTION_VALUE": 3,
"BIN_FINANCIAL_HOME_EQUITY_CREDIT": 3,
"BIN_CONSUMER_REASONER_RECOMMENDATION": 3,
"MUL_HEALTHCARE_WHITE_BLOOD_CELLS_WBC": 3,
"MUL_SOCIAL_DPBEDIA": 3,
"REG_PROFESSIONAL_CPS88_WAGES": 3,
"BIN_SOCIAL_POLITICS_US_CONGRESS_VOTES": 3,
"BIN_GENETICS_OVA_BREAST": 3,
"REG_CONSUMER_BLACK_FRIDAY": 3,
"REG_CONSUMER_AVOCADO_SALES": 3,
"BIN_SCIENCE_CLIMATE_MODEL_SIMULATION_CRASHES": 3,
"BIN_CONSUMER_CHURN_TELCO_NIGERIA": 3,
"BIN_SOCIAL_HATE_SPEECH_DATASET_DYNAMICALLY_GENERATED": 3,
"REG_ANONYM_GEOGRAPHICAL_ORIGIN_OF_MUSIC": 3,
"MUL_HEALTHCARE_CMC_CONTRACEPTIVE_WOMEN": 3,
"BIN_HEALTHCARE_DIABETES_CDC_DEMOGRAPHICS": 3,
"REG_NATURE_POLLEN_LUXEMBOURG": 3,
"REG_FOOD_WHISKY_SCOTCH_REVIEWS": 3,
"BIN_HEALTHCARE_DIABETES_RISK_FACTORS": 3,
"BIN_FINANCIAL_CREDIT_FICO_HELOC": 3,
"REG_FINANCIAL_INSURANCE_PREMIUM_DATA": 3,
"MUL_COMPUTERS_ROBOT_WALL_NAVIGATION": 3,
"BIN_PROFESSIONAL_EMPLOYEE_TURNOVER_TECHCO": 3,
"BIN_SPORTS_CHESS_KR_VS_KP": 3,
"MUL_HEALTHCARE_OBESITY_LEVELS": 3,
"MUL_ANONYM_CNAE": 3,
"BIN_HEALTHCARE_KIDNEY_CHRONIC_DISEASE": 3,
"BIN_INDUSTRIAL_CYLINDER_BANDING_PRINTING": 3,
"MUL_ANONYM_MFEAT_FACTORS": 3,
"REG_TRANSPORTATION_ZURICH_PUBLIC_TRANSPORT_DELAY": 4,
"MUL_FINANCIAL_CREDIT_SCORE_HZL": 4,
"MUL_HEALTHCARE_HEART_ARRHYTMIA": 4,
"MUL_ANONYM_ISOLET_LETTER_SPEECH_RECOGNITION": 4,
"BIN_HEALTHCARE_ALZHEIMER_HANDWRITE_DARWIN": 4,
"MUL_PROFESSIONAL_STUDENT_PERFORMANCE_ENTRANCE_EXAM": 4,
"BIN_FINANCIAL_HOME_CREDIT_DEFAULT_RISK": 4,
"BIN_ANONYM_SATELLITE": 4,
"REG_ANONYM_FAT_MEAT_TECATOR": 4,
"MUL_SPORTS_POKER_HAND": 4,
"BIN_FINANCIAL_CREDIT_GIVE_ME_SOME": 4,
"BIN_HEALTHCARE_INSURANCE_LEAD_PREDICTION": 4,
"BIN_ANONYM_RICARDO": 4,
"BIN_SCIENCE_WATER_TREATMENT": 4,
"REG_ANONYM_MERCEDES_BENZ_GREENER_MANUFACTURING": 4,
"BIN_CONSUMER_CHURN_TELEPHONY": 4,
"BIN_FINANCIAL_LOAN_STATUS_PREDICTION": 4,
"MUL_ANONYM_AMAZON_COMMERCE_REVIEWS": 4,
"MUL_COMPUTERS_IMAGE_GTSRB_GERMAN_TRAFFIC_SIGN": 4,
"REG_FOOD_WINE_JUDGEMENT_SENSORY": 4,
"REG_COMPUTERS_ALGO_RUNTIME_MIP2016": 4,
"BIN_ANONYM_BIORESPONSE": 4,
"BIN_SCIENCE_PARTICLE_MINIBOONE": 4,
"MUL_ANONYM_POL": 4,
"BIN_HEALTHCARE_BREAST_CANCER_MAMMOGRAM": 4,
"REG_HEALTHCARE_POLLUTION_LA_MORTALITY_RMFTSA": 4,
"BIN_CONSUMER_CHURN_TELCO_SOUTH_ASIA": 4,
"REG_ANONYM_HOUSE_16H": 4,
"REG_SPORTS_FIFA_PLAYERS_STATS": 4,
"REG_NATURE_POLLEN_GRAINS_SYNTH": 4,
"BIN_SCIENCE_QSAR_BIODEG": 4,
"BIN_ANONYM_MONKS_PROBLEM_2": 4,
"BIN_HEALTHCARE_CARDIOVASCULAR_DISEASE": 4,
"BIN_SOCIAL_GAMING_LEAGUE_OF_LEGENDS_DIAMOND": 4,
"REG_SCIENCE_SIMULATION_FRIED_FORMULA_PREDICT": 4,
"BIN_ANONYM_APS_FAILURE": 4,
"BIN_FINANCIAL_CREDIT_RISK_DEFAULT": 4,
"MUL_FINANCIAL_CORPORATE_CREDIT_RATING_CAT": 4,
"REG_TRANSPORTATION_SEOUL_BIKE_SHARING_DEMAND": 4,
"BIN_NATURE_PHONEME_SOUNDS": 4,
"MUL_HEALTHCARE_TEETH_DMFT_ANALCATDATA": 4,
"REG_ANONYM_YOLANDA": 4,
"BIN_SCIENCE_EYE_MOVEMENT": 4,
"REG_SCIENCE_SIMULATION_MV": 4,
"REG_SCIENCE_CONCRETE_COMPRESSIVE_STRENGTH": 4,
"MUL_ANONYM_FOURIER_IMAGE_COEFFICIENT": 4,
"REG_HOUSES_BRAZILIAN_HOUSES": 4,
"BIN_HEALTHCARE_LUNG_MALE_ARSENIC": 4,
"REG_TRANSPORTATION_NYC_BIKE_TRIP_DURATION": 4,
"BIN_ANONYM_MADELINE": 4,
"REG_CONSUMER_ONLINE_NEWS_POPULARITY": 4,
"MUL_COMPUTERS_IMAGE_INDIAN_PINES": 4,
"BIN_SCIENCE_ANTIGEN_AVIDA_HIL6": 4,
"MUL_COMPUTERS_META_STREAM_INTERVALS": 4,
"MUL_SPORTS_CONNECT4_GAME": 4,
"MUL_HEALTHCARE_LYMPH_YUGOSLAVIA": 4,
"BIN_SCIENCE_SIMULATION_HILL_VALLEY": 4,
"BIN_CONSUMER_NEWSPAPER_CHURN": 4,
"BIN_FINANCIAL_CC_FRAUD": 4,
"MUL_SPORTS_WALKING_ACTIVITY": 4,
"REG_COMPUTERS_GAMING_FRAMES_FPS_BENCHMARK": 4,
"REG_NATURE_QUAKE_RICHTER": 4,
"MUL_ANONYM_VOLKERT": 4,
"BIN_COMPUTERS_IMAGE_SKIN_SEGMENTATION": 4,
"BIN_COMPUTERS_IMAGE_BANK_NOTE_AUTHENTICATION": 4,
"REG_SPORTS_NBA_ALL_STAR": 4,
}


def get_tabstar_version(pretrain_dataset_or_path: Optional[str] = None) -> str:
    if pretrain_dataset_or_path is None:
        return "alana89/TabSTAR"
    if pretrain_dataset_or_path.startswith(("BIN_", "REG_", "MUL_")):
        tabstar_version = _get_tabstar_version_from_dataset(pretrain_dataset=pretrain_dataset_or_path)
        return f"alana89/{tabstar_version}"
    if isinstance(pretrain_dataset_or_path, str):
        return pretrain_dataset_or_path
    raise ValueError(f"Unknown pretrain_dataset_or_path: {pretrain_dataset_or_path}")


def _get_tabstar_version_from_dataset(pretrain_dataset: str) -> str:
    text_fold = TEXT2FOLD.get(pretrain_dataset)
    if text_fold is not None:
        return f"TabSTAR-paper-version-fold-k{text_fold}"
    pretrain_fold = PRETRAIN2FOLD.get(pretrain_dataset)
    if pretrain_fold is not None:
        return f"TabSTAR-eval-320-version-fold-k{pretrain_fold}"
    print(f"🤷‍♂️ Unknown dataset: {pretrain_dataset}, falling back to TabSTAR")
    return "TabSTAR"