""" Fine-tuning GLiNER large-v2.1 for data mention extraction. Entity types: DATASET - Named datasets (DHS dataset, Census microdata, LSMS-ISA panel dataset) SURVEY - Named survey programs (Demographic and Health Survey, MICS, LSMS) DATABASE - Named databases/repositories (World Development Indicators, IHME, FAOSTAT) GEOCODED_DATA - Geocoded/spatial data (GPS coordinates, geocoded cluster coordinates) VAGUE_MENTION - Vague/informal references (a survey from Ghana, a national household survey) HOW TO RUN (HF Jobs): hf_jobs run train_gliner.py \ --hardware t4-small \ --timeout 3h \ --dependencies "gliner==0.2.26 torch trackio" \ --trackio_space_id rafmacalaba/ml-intern-gliner01 \ --trackio_project gliner-data-mentions HOW TO RUN (local): pip install gliner==0.2.26 torch trackio python train_gliner.py python train_gliner.py --smoke_test # 2-epoch sanity check """ import json import os import math import random import torch from torch.utils.data import Dataset from gliner import GLiNER from gliner.training import Trainer, TrainingArguments # -------------------------------------------------------------------------- # Inline training data (121 examples, 175 spans, 5 entity types) # All data is embedded directly so the job is fully self-contained. # -------------------------------------------------------------------------- TRAIN_DATA_JSON = """ [ {"tokenized_text": ["The", "Demographic", "and", "Health", "Survey", "was", "conducted", "in", "47", "countries", "across", "sub-Saharan", "Africa."], "ner": [[1, 5, "SURVEY"]]}, {"tokenized_text": ["We", "used", "the", "2018", "Demographic", "and", "Health", "Survey", "for", "Ghana", "to", "estimate", "child", "mortality", "rates."], "ner": [[4, 8, "SURVEY"]]}, {"tokenized_text": ["DHS", "data", "collected", "from", "household", "interviews", "provided", "birth", "history", "records", "for", "all", "women", "aged", "15-49."], "ner": [[0, 1, "SURVEY"]]}, {"tokenized_text": ["The", "DHS", "Program", "provides", "nationally", "representative", "data", "on", "population", "health", "and", "nutrition."], "ner": [[1, 3, "SURVEY"]]}, {"tokenized_text": ["Data", "from", "the", "2016", "Ghana", "DHS", "were", "linked", "to", "GPS", "coordinates", "of", "enumeration", "areas."], "ner": [[4, 6, "SURVEY"], [9, 11, "GEOCODED_DATA"]]}, {"tokenized_text": ["The", "Multiple", "Indicator", "Cluster", "Survey", "is", "conducted", "by", "UNICEF", "across", "low-income", "countries."], "ner": [[1, 5, "SURVEY"]]}, {"tokenized_text": ["MICS", "round", "6", "data", "for", "15", "West", "African", "countries", "were", "used", "to", "assess", "vaccination", "coverage."], "ner": [[0, 1, "SURVEY"]]}, {"tokenized_text": ["The", "Living", "Standards", "Measurement", "Study", "surveys", "track", "household", "welfare", "over", "time."], "ner": [[1, 5, "SURVEY"]]}, {"tokenized_text": ["LSMS-ISA", "data", "were", "used", "to", "study", "agricultural", "productivity", "in", "Ethiopia."], "ner": [[0, 1, "SURVEY"]]}, {"tokenized_text": ["We", "linked", "LSMS", "household", "locations", "with", "satellite-derived", "rainfall", "estimates."], "ner": [[2, 3, "SURVEY"]]}, {"tokenized_text": ["Afrobarometer", "survey", "data", "measured", "democratic", "attitudes", "in", "36", "African", "countries."], "ner": [[0, 1, "SURVEY"]]}, {"tokenized_text": ["Round", "8", "of", "the", "Afrobarometer", "was", "collected", "in", "2019-2021."], "ner": [[4, 5, "SURVEY"]]}, {"tokenized_text": ["Ethiopia", "Demographic", "and", "Health", "Survey", "2016", "data", "were", "analyzed", "for", "stunting", "prevalence."], "ner": [[0, 6, "SURVEY"]]}, {"tokenized_text": ["Tanzania", "DHS", "2015-16", "provided", "cluster-level", "data", "for", "8593", "households."], "ner": [[0, 3, "SURVEY"]]}, {"tokenized_text": ["Data", "from", "the", "Nigeria", "Demographic", "and", "Health", "Survey", "were", "analyzed."], "ner": [[3, 8, "SURVEY"]]}, {"tokenized_text": ["The", "South", "Africa", "Demographic", "and", "Health", "Survey", "used", "a", "stratified", "two-stage", "cluster", "design."], "ner": [[1, 7, "SURVEY"]]}, {"tokenized_text": ["Nepal", "DHS", "2022", "included", "GPS", "coordinates", "of", "surveyed", "households."], "ner": [[0, 3, "SURVEY"], [4, 6, "GEOCODED_DATA"]]}, {"tokenized_text": ["The", "National", "Health", "and", "Demographic", "Survey", "was", "conducted", "every", "five", "years."], "ner": [[1, 6, "SURVEY"]]}, {"tokenized_text": ["The", "Malawi", "Integrated", "Household", "Survey", "provides", "longitudinal", "welfare", "data."], "ner": [[1, 5, "SURVEY"]]}, {"tokenized_text": ["Kenya", "DHS", "2022", "provided", "geocoded", "cluster", "locations", "with", "a", "2", "km", "displacement."], "ner": [[0, 3, "SURVEY"], [4, 7, "GEOCODED_DATA"]]}, {"tokenized_text": ["WHO", "SARA", "survey", "data", "were", "used", "to", "assess", "health", "facility", "readiness."], "ner": [[0, 3, "SURVEY"]]}, {"tokenized_text": ["The", "Demographic", "and", "Health", "Surveys", "program", "has", "collected", "data", "since", "1984."], "ner": [[1, 5, "SURVEY"]]}, {"tokenized_text": ["We", "analyzed", "data", "from", "the", "India", "NFHS-5", "conducted", "in", "2019-2021."], "ner": [[5, 7, "SURVEY"]]}, {"tokenized_text": ["The", "Malaria", "Indicator", "Survey", "was", "implemented", "in", "12", "high-burden", "countries."], "ner": [[1, 4, "SURVEY"]]}, {"tokenized_text": ["USAID", "supported", "the", "2019", "Demographic", "and", "Health", "Survey", "in", "Mozambique."], "ner": [[4, 8, "SURVEY"]]}, {"tokenized_text": ["Bangladesh", "Demographic", "and", "Health", "Survey", "2017-18", "covered", "20127", "households."], "ner": [[0, 6, "SURVEY"]]}, {"tokenized_text": ["The", "Indonesia", "Family", "Life", "Survey", "tracked", "more", "than", "7000", "households."], "ner": [[1, 5, "SURVEY"]]}, {"tokenized_text": ["Household", "data", "from", "the", "Pan", "African", "Survey", "were", "used", "to", "compare", "outcomes."], "ner": [[4, 7, "SURVEY"]]}, {"tokenized_text": ["The", "Integrated", "Household", "Living", "Conditions", "Survey", "from", "Rwanda", "was", "accessed."], "ner": [[1, 6, "SURVEY"]]}, {"tokenized_text": ["MICS", "data", "for", "Somalia", "were", "used", "to", "calculate", "child", "stunting", "rates."], "ner": [[0, 1, "SURVEY"]]}, {"tokenized_text": ["National", "population", "census", "data", "were", "obtained", "from", "the", "Ghana", "Statistical", "Service."], "ner": [[0, 4, "DATASET"]]}, {"tokenized_text": ["The", "2010", "Population", "and", "Housing", "Census", "covered", "all", "enumeration", "areas", "in", "Ghana."], "ner": [[1, 6, "DATASET"]]}, {"tokenized_text": ["Census", "microdata", "at", "the", "household", "level", "were", "linked", "with", "satellite", "imagery."], "ner": [[0, 2, "DATASET"]]}, {"tokenized_text": ["We", "used", "the", "national", "census", "to", "derive", "population", "denominators", "for", "each", "district."], "ner": [[3, 5, "DATASET"]]}, {"tokenized_text": ["The", "2000", "census", "provided", "the", "sampling", "frame", "for", "subsequent", "household", "surveys."], "ner": [[1, 3, "DATASET"]]}, {"tokenized_text": ["Census", "data", "from", "the", "2015", "round", "were", "used", "to", "estimate", "poverty", "rates."], "ner": [[0, 2, "DATASET"]]}, {"tokenized_text": ["Full-count", "census", "microdata", "for", "Nigeria", "were", "geo-referenced", "at", "enumeration", "area", "level."], "ner": [[0, 3, "DATASET"], [6, 7, "GEOCODED_DATA"]]}, {"tokenized_text": ["The", "DHS", "dataset", "contains", "individual-level", "records", "for", "women", "of", "reproductive", "age."], "ner": [[1, 3, "DATASET"]]}, {"tokenized_text": ["We", "used", "MICS", "household", "survey", "data", "for", "water", "sanitation", "and", "hygiene", "analysis."], "ner": [[2, 6, "DATASET"], [2, 3, "SURVEY"]]}, {"tokenized_text": ["The", "LSMS-ISA", "panel", "dataset", "spans", "six", "survey", "rounds", "between", "2010", "and", "2019."], "ner": [[1, 4, "DATASET"]]}, {"tokenized_text": ["Population", "census", "microdata", "from", "10", "countries", "were", "harmonized", "by", "IPUMS."], "ner": [[0, 3, "DATASET"]]}, {"tokenized_text": ["The", "2017", "Sudan", "Multiple", "Indicator", "Cluster", "Survey", "dataset", "was", "used."], "ner": [[1, 8, "DATASET"], [2, 7, "SURVEY"]]}, {"tokenized_text": ["We", "merged", "the", "2018", "DHS", "dataset", "with", "geocoded", "household", "locations."], "ner": [[3, 6, "DATASET"], [7, 9, "GEOCODED_DATA"]]}, {"tokenized_text": ["The", "study", "combined", "the", "Living", "Standards", "Measurement", "Study", "with", "census", "records."], "ner": [[4, 8, "SURVEY"], [9, 11, "DATASET"]]}, {"tokenized_text": ["Our", "analysis", "uses", "DHS", "census", "microdata", "and", "administrative", "health", "records."], "ner": [[3, 4, "SURVEY"], [4, 6, "DATASET"], [7, 10, "DATABASE"]]}, {"tokenized_text": ["World", "Development", "Indicators", "database", "was", "accessed", "to", "retrieve", "GDP", "per", "capita", "data."], "ner": [[0, 4, "DATABASE"]]}, {"tokenized_text": ["Data", "from", "the", "World", "Bank", "Open", "Data", "portal", "were", "downloaded", "in", "July", "2023."], "ner": [[3, 8, "DATABASE"]]}, {"tokenized_text": ["The", "PovcalNet", "database", "provides", "harmonized", "consumption", "and", "poverty", "estimates."], "ner": [[1, 3, "DATABASE"]]}, {"tokenized_text": ["Global", "Burden", "of", "Disease", "estimates", "from", "IHME", "were", "used", "in", "this", "analysis."], "ner": [[0, 4, "DATABASE"], [6, 7, "DATABASE"]]}, {"tokenized_text": ["GBD", "2019", "data", "provided", "country-level", "DALY", "estimates."], "ner": [[0, 3, "DATABASE"]]}, {"tokenized_text": ["IPUMS-DHS", "integrates", "harmonized", "variables", "from", "the", "Demographic", "and", "Health", "Surveys."], "ner": [[0, 1, "DATABASE"], [6, 10, "SURVEY"]]}, {"tokenized_text": ["We", "used", "IPUMS", "International", "census", "microdata", "for", "Kenya", "Uganda", "and", "Tanzania."], "ner": [[2, 4, "DATABASE"]]}, {"tokenized_text": ["Ministry", "of", "Health", "routine", "health", "information", "system", "records", "were", "linked", "with", "survey", "data."], "ner": [[3, 8, "DATABASE"]]}, {"tokenized_text": ["We", "combined", "birth", "registration", "data", "with", "household", "survey", "responses."], "ner": [[2, 5, "DATABASE"]]}, {"tokenized_text": ["Administrative", "records", "from", "district", "health", "offices", "were", "used", "as", "a", "secondary", "data", "source."], "ner": [[0, 2, "DATABASE"]]}, {"tokenized_text": ["The", "Malnutrition", "database", "maintained", "by", "WHO", "was", "queried", "for", "country-level", "estimates."], "ner": [[1, 3, "DATABASE"]]}, {"tokenized_text": ["Global", "Health", "Observatory", "data", "repository", "was", "used", "to", "extract", "neonatal", "mortality", "rates."], "ner": [[0, 5, "DATABASE"]]}, {"tokenized_text": ["The", "UNAIDS", "data", "repository", "provided", "country", "estimates", "of", "HIV", "prevalence."], "ner": [[1, 4, "DATABASE"]]}, {"tokenized_text": ["We", "accessed", "the", "UNICEF", "MICS", "data", "portal", "to", "download", "cleaned", "survey", "files."], "ner": [[3, 7, "DATABASE"], [4, 5, "SURVEY"]]}, {"tokenized_text": ["FAOSTAT", "database", "provided", "data", "on", "food", "production", "and", "agricultural", "land", "use."], "ner": [[0, 2, "DATABASE"]]}, {"tokenized_text": ["The", "IMF", "World", "Economic", "Outlook", "database", "was", "accessed", "for", "fiscal", "data."], "ner": [[1, 6, "DATABASE"]]}, {"tokenized_text": ["We", "used", "the", "Human", "Mortality", "Database", "to", "compute", "life", "expectancy", "at", "birth."], "ner": [[3, 6, "DATABASE"]]}, {"tokenized_text": ["The", "GRID3", "database", "includes", "geocoded", "administrative", "boundaries", "and", "health", "facility", "locations."], "ner": [[1, 3, "DATABASE"], [4, 7, "GEOCODED_DATA"]]}, {"tokenized_text": ["Routine", "health", "data", "from", "the", "District", "Health", "Information", "System", "were", "extracted."], "ner": [[5, 9, "DATABASE"]]}, {"tokenized_text": ["WorldPop", "gridded", "population", "data", "were", "used", "to", "weight", "survey", "estimates."], "ner": [[0, 4, "DATABASE"]]}, {"tokenized_text": ["We", "downloaded", "OpenStreetMap", "data", "to", "compute", "road", "distance", "metrics."], "ner": [[2, 4, "DATABASE"]]}, {"tokenized_text": ["The", "Global", "Health", "Security", "Index", "database", "was", "consulted", "for", "preparedness", "scores."], "ner": [[1, 6, "DATABASE"]]}, {"tokenized_text": ["The", "ACLED", "conflict", "event", "database", "was", "used", "to", "measure", "political", "instability."], "ner": [[1, 5, "DATABASE"]]}, {"tokenized_text": ["We", "used", "the", "USAID", "FoodNet", "database", "to", "assess", "nutritional", "status."], "ner": [[3, 6, "DATABASE"]]}, {"tokenized_text": ["Geocoded", "household", "survey", "data", "were", "overlaid", "with", "flood", "inundation", "maps."], "ner": [[0, 4, "GEOCODED_DATA"]]}, {"tokenized_text": ["We", "used", "GPS-tagged", "household", "locations", "from", "the", "2014", "Ethiopia", "DHS."], "ner": [[2, 5, "GEOCODED_DATA"], [7, 10, "SURVEY"]]}, {"tokenized_text": ["Spatial", "household", "data", "were", "linked", "to", "land", "cover", "classifications", "from", "MODIS."], "ner": [[0, 3, "GEOCODED_DATA"]]}, {"tokenized_text": ["The", "study", "linked", "geocoded", "enumeration", "areas", "to", "nighttime", "light", "intensity", "values."], "ner": [[3, 6, "GEOCODED_DATA"]]}, {"tokenized_text": ["Geolocated", "survey", "clusters", "were", "matched", "to", "district-level", "infrastructure", "access", "scores."], "ner": [[0, 3, "GEOCODED_DATA"]]}, {"tokenized_text": ["GPS", "coordinates", "from", "household", "surveys", "were", "used", "to", "assign", "climate", "zone", "classifications."], "ner": [[0, 2, "GEOCODED_DATA"]]}, {"tokenized_text": ["Geocoded", "hospital", "locations", "were", "linked", "to", "DHS", "cluster", "coordinates."], "ner": [[0, 3, "GEOCODED_DATA"], [6, 9, "GEOCODED_DATA"]]}, {"tokenized_text": ["DHS", "GPS", "data", "include", "geocoded", "cluster", "locations", "displaced", "10", "km", "in", "rural", "areas."], "ner": [[0, 3, "GEOCODED_DATA"], [4, 7, "GEOCODED_DATA"]]}, {"tokenized_text": ["We", "extracted", "geocoded", "DHS", "cluster", "coordinates", "to", "map", "malaria", "prevalence."], "ner": [[2, 6, "GEOCODED_DATA"]]}, {"tokenized_text": ["The", "Demographic", "and", "Health", "Survey", "GPS", "dataset", "provides", "displaced", "geographic", "coordinates."], "ner": [[1, 7, "GEOCODED_DATA"], [1, 5, "SURVEY"]]}, {"tokenized_text": ["LSMS-ISA", "geocoded", "panel", "data", "were", "used", "to", "study", "agricultural", "productivity."], "ner": [[0, 4, "GEOCODED_DATA"], [0, 1, "SURVEY"]]}, {"tokenized_text": ["We", "linked", "DHS", "cluster", "centroids", "with", "500m", "resolution", "land", "surface", "temperature", "rasters."], "ner": [[2, 5, "GEOCODED_DATA"]]}, {"tokenized_text": ["Household", "GPS", "data", "were", "used", "to", "assign", "malaria", "transmission", "intensity", "zones."], "ner": [[0, 3, "GEOCODED_DATA"]]}, {"tokenized_text": ["Geo-referenced", "health", "facility", "data", "from", "the", "WHO", "SARA", "survey", "were", "used."], "ner": [[0, 4, "GEOCODED_DATA"], [6, 9, "SURVEY"]]}, {"tokenized_text": ["Tanzania", "DHS", "2015-16", "provided", "geocoded", "cluster-level", "data", "for", "8593", "households."], "ner": [[0, 3, "SURVEY"], [4, 7, "GEOCODED_DATA"]]}, {"tokenized_text": ["Data", "from", "the", "Nigeria", "Demographic", "and", "Health", "Survey", "were", "linked", "to", "georeferenced", "malaria", "maps."], "ner": [[3, 8, "SURVEY"], [11, 14, "GEOCODED_DATA"]]}, {"tokenized_text": ["Nepal", "DHS", "2022", "included", "GPS", "coordinates", "of", "surveyed", "households."], "ner": [[0, 3, "SURVEY"], [4, 6, "GEOCODED_DATA"]]}, {"tokenized_text": ["GPS", "coordinates", "from", "the", "2019", "Kenya", "DHS", "were", "matched", "to", "GRID3", "building", "footprints."], "ner": [[0, 2, "GEOCODED_DATA"], [4, 7, "SURVEY"]]}, {"tokenized_text": ["Geocoded", "data", "from", "the", "LSMS-ISA", "were", "merged", "with", "soil", "quality", "maps."], "ner": [[0, 2, "GEOCODED_DATA"], [4, 5, "SURVEY"]]}, {"tokenized_text": ["The", "study", "used", "geostatistical", "data", "from", "the", "Malaria", "Atlas", "Project."], "ner": [[3, 5, "GEOCODED_DATA"]]}, {"tokenized_text": ["We", "overlaid", "geocoded", "health", "facility", "locations", "with", "travel", "time", "rasters."], "ner": [[2, 6, "GEOCODED_DATA"]]}, {"tokenized_text": ["Spatially", "referenced", "household", "data", "were", "linked", "to", "conflict", "event", "records."], "ner": [[0, 4, "GEOCODED_DATA"]]}, {"tokenized_text": ["We", "used", "GPS", "location", "data", "from", "survey", "clusters", "to", "compute", "distance", "to", "markets."], "ner": [[2, 5, "GEOCODED_DATA"]]}, {"tokenized_text": ["Geo-coded", "wealth", "index", "scores", "were", "mapped", "across", "200", "districts."], "ner": [[0, 4, "GEOCODED_DATA"]]}, {"tokenized_text": ["A", "survey", "from", "Ghana", "collected", "data", "on", "child", "nutrition", "and", "household", "wealth."], "ner": [[0, 4, "VAGUE_MENTION"]]}, {"tokenized_text": ["The", "study", "relied", "on", "a", "nationally", "representative", "household", "survey."], "ner": [[4, 9, "VAGUE_MENTION"]]}, {"tokenized_text": ["We", "used", "data", "from", "a", "cross-sectional", "survey", "in", "rural", "Ethiopia."], "ner": [[4, 10, "VAGUE_MENTION"]]}, {"tokenized_text": ["Administrative", "data", "from", "local", "government", "offices", "were", "used", "in", "the", "regression", "model."], "ner": [[0, 6, "VAGUE_MENTION"]]}, {"tokenized_text": ["Secondary", "analysis", "of", "a", "household", "survey", "dataset", "from", "2015", "was", "conducted."], "ner": [[3, 9, "VAGUE_MENTION"]]}, {"tokenized_text": ["Several", "household", "surveys", "from", "sub-Saharan", "Africa", "were", "pooled", "for", "the", "meta-analysis."], "ner": [[0, 6, "VAGUE_MENTION"]]}, {"tokenized_text": ["The", "analysis", "drew", "on", "population-based", "surveys", "from", "three", "East", "African", "countries."], "ner": [[4, 11, "VAGUE_MENTION"]]}, {"tokenized_text": ["Data", "from", "a", "regional", "household", "survey", "conducted", "in", "2017", "were", "used."], "ner": [[2, 9, "VAGUE_MENTION"]]}, {"tokenized_text": ["A", "large-scale", "national", "survey", "in", "Kenya", "collected", "data", "on", "reproductive", "health."], "ner": [[0, 6, "VAGUE_MENTION"]]}, {"tokenized_text": ["A", "household", "panel", "study", "from", "Zimbabwe", "tracked", "changes", "in", "food", "security", "over", "time."], "ner": [[0, 6, "VAGUE_MENTION"]]}, {"tokenized_text": ["National", "survey", "data", "from", "Bangladesh", "were", "used", "as", "the", "primary", "data", "source."], "ner": [[0, 5, "VAGUE_MENTION"]]}, {"tokenized_text": ["The", "study", "used", "administrative", "health", "data", "from", "five", "districts", "in", "Uganda."], "ner": [[3, 11, "VAGUE_MENTION"]]}, {"tokenized_text": ["Data", "from", "a", "longitudinal", "cohort", "study", "in", "South", "Africa", "were", "used."], "ner": [[2, 9, "VAGUE_MENTION"]]}, {"tokenized_text": ["The", "research", "relied", "on", "a", "cross-national", "survey", "of", "household", "incomes."], "ner": [[4, 10, "VAGUE_MENTION"]]}, {"tokenized_text": ["LSMS", "data", "the", "2010", "census", "and", "World", "Development", "Indicators", "were", "used", "in", "the", "model."], "ner": [[0, 2, "DATASET"], [3, 5, "DATASET"], [6, 9, "DATABASE"]]}, {"tokenized_text": ["GPS", "coordinates", "of", "DHS", "clusters", "were", "matched", "with", "WorldPop", "gridded", "population", "data."], "ner": [[0, 2, "GEOCODED_DATA"], [3, 4, "SURVEY"], [8, 12, "DATABASE"]]}, {"tokenized_text": ["We", "merged", "the", "2018", "DHS", "dataset", "with", "geocoded", "household", "locations", "and", "census", "records."], "ner": [[3, 6, "DATASET"], [7, 10, "GEOCODED_DATA"], [11, 13, "DATASET"]]}, {"tokenized_text": ["Data", "from", "the", "Demographic", "and", "Health", "Survey", "a", "survey", "from", "Nigeria", "and", "administrative", "databases", "were", "pooled."], "ner": [[3, 7, "SURVEY"], [7, 11, "VAGUE_MENTION"], [12, 14, "DATABASE"]]}, {"tokenized_text": ["IPUMS-DHS", "integrates", "harmonized", "variables", "from", "the", "Demographic", "and", "Health", "Surveys."], "ner": [[0, 1, "DATABASE"], [6, 10, "SURVEY"]]}, {"tokenized_text": ["We", "accessed", "the", "UNICEF", "MICS", "data", "portal", "to", "download", "the", "MICS", "dataset."], "ner": [[3, 7, "DATABASE"], [10, 12, "DATASET"], [4, 5, "SURVEY"]]}, {"tokenized_text": ["Three", "datasets", "were", "used:", "the", "2016", "MICS", "for", "Ethiopia", "the", "2017", "NLSS", "for", "Nepal", "and", "a", "regional", "survey", "from", "2018."], "ner": [[5, 9, "DATASET"], [6, 7, "SURVEY"], [10, 14, "DATASET"], [15, 20, "VAGUE_MENTION"]]}, {"tokenized_text": ["Geocoded", "data", "from", "the", "LSMS-ISA", "were", "merged", "with", "soil", "quality", "maps", "and", "DHS", "cluster", "locations."], "ner": [[0, 2, "GEOCODED_DATA"], [4, 5, "SURVEY"], [12, 15, "GEOCODED_DATA"]]}, {"tokenized_text": ["We", "used", "data", "from", "national", "censuses", "and", "a", "household", "survey", "from", "Malawi."], "ner": [[4, 6, "DATASET"], [7, 12, "VAGUE_MENTION"]]}, {"tokenized_text": ["DHS", "cluster", "centroids", "were", "overlaid", "with", "GRID3", "settlement", "layer", "data."], "ner": [[0, 3, "GEOCODED_DATA"], [0, 1, "SURVEY"], [6, 10, "DATABASE"]]}, {"tokenized_text": ["The", "2014", "Ethiopia", "DHS", "GPS", "dataset", "was", "used", "to", "extract", "geocoded", "household", "cluster", "positions."], "ner": [[1, 6, "GEOCODED_DATA"], [1, 4, "SURVEY"], [10, 14, "GEOCODED_DATA"]]}, {"tokenized_text": ["We", "combined", "DHS", "survey", "data", "with", "geospatial", "layers", "from", "WorldPop", "and", "ACLED."], "ner": [[2, 5, "DATASET"], [2, 3, "SURVEY"], [9, 10, "DATABASE"], [11, 12, "DATABASE"]]}, {"tokenized_text": ["Administrative", "records", "birth", "registration", "data", "and", "a", "household", "survey", "from", "Niger", "were", "analyzed."], "ner": [[0, 2, "DATABASE"], [2, 5, "DATABASE"], [6, 11, "VAGUE_MENTION"]]}, {"tokenized_text": ["Data", "from", "GBD", "2019", "were", "combined", "with", "georeferenced", "DHS", "cluster", "data", "and", "ACLED", "conflict", "records."], "ner": [[2, 4, "DATABASE"], [7, 11, "GEOCODED_DATA"], [8, 9, "SURVEY"], [12, 15, "DATABASE"]]} ] """ ENTITY_LABELS = ["DATASET", "SURVEY", "DATABASE", "GEOCODED_DATA", "VAGUE_MENTION"] class GLiNERDataset(Dataset): def __init__(self, data): self.data = data def __len__(self): return len(self.data) def __getitem__(self, idx): item = dict(self.data[idx]) item["ner"] = [tuple(s) for s in item["ner"]] return item def main(): import argparse parser = argparse.ArgumentParser() parser.add_argument("--smoke_test", action="store_true") parser.add_argument("--base_model", default="urchade/gliner_large-v2.1") parser.add_argument("--output_dir", default="./gliner-data-mentions") parser.add_argument("--hub_model_id", default="rafmacalaba/gliner-large-data-mentions") parser.add_argument("--num_epochs", type=int, default=25) parser.add_argument("--batch_size", type=int, default=8) parser.add_argument("--lr", type=float, default=1e-5) parser.add_argument("--others_lr", type=float, default=5e-5) args = parser.parse_args() all_data = json.loads(TRAIN_DATA_JSON) random.seed(42) random.shuffle(all_data) split_idx = int(len(all_data) * 0.85) train_data = all_data[:split_idx] eval_data = all_data[split_idx:] if args.smoke_test: print("SMOKE TEST MODE") train_data = train_data[:20] eval_data = eval_data[:5] args.num_epochs = 2 args.batch_size = 4 print(f"Train: {len(train_data)}, Eval: {len(eval_data)}") print(f"Loading {args.base_model} ...") model = GLiNER.from_pretrained(args.base_model) collator = model._create_data_collator() total_steps = math.ceil(len(train_data) / args.batch_size) * args.num_epochs warmup_steps = max(1, int(total_steps * 0.08)) training_args = TrainingArguments( output_dir=args.output_dir, num_train_epochs=args.num_epochs, per_device_train_batch_size=args.batch_size, gradient_accumulation_steps=1, learning_rate=args.lr, others_lr=args.others_lr, weight_decay=0.01, warmup_steps=warmup_steps, lr_scheduler_type="cosine", save_strategy="epoch", save_total_limit=3, load_best_model_at_end=False, logging_steps=5, logging_first_step=True, logging_strategy="steps", disable_tqdm=True, fp16=torch.cuda.is_available(), report_to="trackio", run_name="gliner-large-data-mentions-v1", project=os.environ.get("TRACKIO_PROJECT", "gliner-data-mentions"), trackio_space_id=os.environ.get("TRACKIO_SPACE_ID", "rafmacalaba/ml-intern-gliner01"), push_to_hub=True, hub_model_id=args.hub_model_id, hub_strategy="end", loss_reduction="sum", negatives=1.0, focal_loss_alpha=-1, focal_loss_gamma=0, masking="global", ) trainer = Trainer( model=model, args=training_args, train_dataset=GLiNERDataset(train_data), data_collator=collator, ) print("=== Starting training ===") trainer.train() print("Training complete.") save_path = os.path.join(args.output_dir, "final") model.save_pretrained(save_path) print(f"Saved to {save_path}") hf_token = os.environ.get("HF_TOKEN") if hf_token and args.hub_model_id and not args.smoke_test: print(f"Pushing to Hub: {args.hub_model_id}") model.push_to_hub(args.hub_model_id, token=hf_token) print("Hub push successful!") print("\n=== Inference examples ===") test_sentences = [ "The Demographic and Health Survey collected household data in 2018.", "Geocoded DHS cluster coordinates were overlaid with MODIS land cover.", "We used data from a household survey conducted in rural Ghana.", "GBD 2019 data from IHME were combined with WorldPop gridded population data.", "IPUMS-DHS integrates harmonized variables from the Demographic and Health Surveys.", ] loaded = GLiNER.from_pretrained(save_path) for text in test_sentences: ents = loaded.predict_entities(text, labels=ENTITY_LABELS, threshold=0.4) print(f"\n TEXT: {text}") for e in ents: print(f" [{e['label']}] '{e['text']}' ({e['score']:.3f})") if __name__ == "__main__": main()