create_dane.py

# /// script
# requires-python = ">=3.10,<4.0"
# dependencies = [
#     "datasets==3.5.0",
#     "huggingface-hub==0.24.0",
#     "pandas==2.2.0",
#     "requests==2.32.3",
# ]
# ///

"""Create the DaNE-mini NER dataset and upload it to the HF Hub."""

import re
from collections import defaultdict

import pandas as pd
import requests
from datasets import Split
from datasets.arrow_dataset import Dataset
from datasets.dataset_dict import DatasetDict
from huggingface_hub.hf_api import HfApi


def main() -> None:
    """Create the DaNE-mini NER dataset and uploads it to the HF Hub."""
    # Define download URLs
    # TODO: This is the wrong URL; use the alexandra URL instead
    base_url = (
        "https://github.com/UniversalDependencies/UD_Danish-DDT/raw/master/"
        "da_ddt-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    # Download the data
    data = dict(
        train=requests.get(train_url).text.split("\n"),
        val=requests.get(val_url).text.split("\n"),
        test=requests.get(test_url).text.split("\n"),
    )

    # Iterate over the data splits
    dfs = dict()
    for split, lines in data.items():
        # Initialise the records, data dictionary and document
        records = list()
        data_dict: dict[str, list[int | str]] = defaultdict(list)
        doc = ""

        # Iterate over the data for the given split
        for line in lines:
            # If we are at the first line of an entry then extract the document
            if line.startswith("# text = "):
                doc = re.sub("# text = ", "", line)

            # Otherwise, if the line is a comment then ignore it
            elif line.startswith("#"):
                continue

            # Otherwise, if we have reached the end of an entry then store it to the
            # list of records and reset the data dictionary and document
            elif line == "":
                if len(data_dict["tokens"]) > 0:
                    merged_data_dict: dict[str, str | list[int | str]]
                    merged_data_dict = {**data_dict, "doc": doc}
                    records.append(merged_data_dict)
                data_dict = defaultdict(list)
                doc = ""

            # Otherwise we are in the middle of an entry which is not a comment, so
            # we extract the data from the line and store it in the data dictionary
            else:
                data_tup = line.split("\t")
                ner_tag = data_tup[9].replace("name=", "").split("|")[0]
                data_dict["ids"].append(data_tup[0])
                data_dict["tokens"].append(data_tup[1])
                data_dict["ner_tags"].append(ner_tag)

        # Convert the records to a dataframe and store it
        dfs[split] = pd.DataFrame.from_records(records)

    # Merge the dataframes
    df = pd.concat(dfs.values()).reset_index(drop=True)

    # Remove duplicates
    df = df.drop_duplicates().reset_index(drop=True)

    # Create new splits
    val_df = df.sample(n=256, random_state=4242)
    df_filtered = df.loc[~df.index.isin(val_df.index)]
    test_df = df_filtered.sample(n=2048, random_state=4242)
    full_train_df = df_filtered.loc[~df_filtered.index.isin(test_df.index)]
    train_df = full_train_df.sample(n=1024, random_state=4242)

    # Collect datasets in a dataset dictionary
    dataset = DatasetDict(
        train=Dataset.from_pandas(train_df, split=Split.TRAIN),
        val=Dataset.from_pandas(val_df, split=Split.VALIDATION),
        test=Dataset.from_pandas(test_df, split=Split.TEST),
        full_train=Dataset.from_pandas(full_train_df, split="full_train"),  # type: ignore[bad-argument-type]
    )

    # Create dataset ID
    dataset_id = "EuroEval/dane-mini"

    # Remove the dataset from Hugging Face Hub if it already exists
    HfApi().delete_repo(dataset_id, repo_type="dataset", missing_ok=True)

    # Push the dataset to the Hugging Face Hub
    dataset.push_to_hub(dataset_id, private=True)


if __name__ == "__main__":
    main()