Initial upload

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+sp.vocab filter=lfs diff=lfs merge=lfs -text

dataset/o3-mini/20250218/data-00000-of-00001.arrow

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d78eaab378f462e45adbf026f13c3d5b9a289ddea75d399dbbf0c12b0c25c11e
+size 40179808

dataset/o3-mini/20250218/dataset_info.json

@@ -0,0 +1,110 @@
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "text": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "tokens": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "adj": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "adv": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "det": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "enc": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "func": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "misc": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "ner1": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "ner2": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "noun": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "pronoun": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "punct": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "verb": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "wh": {
+      "feature": {
+        "dtype": "string",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

dataset/o3-mini/20250218/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "4a79c58b9023cf85",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

dataset_maker.py

@@ -0,0 +1,420 @@
+from asyncio import Task
+
+from datasets import Dataset, load_dataset
+from openai import AsyncOpenAI
+from traceback import format_exc
+from typing import Union
+import asyncio
+import itertools
+import json
+import logging.config
+import random
+import sentencepiece as spm
+
+from utils import default_logging_config
+
+client = AsyncOpenAI()
+logger = logging.getLogger(__name__)
+
+sp = spm.SentencePieceProcessor()
+sp.LoadFromFile(f"sp.model")
+
+#OPENAI_MODEL = "gpt-4o"
+#OPENAI_MODEL = "o1"
+OPENAI_MODEL = "o3-mini"
+
+
+async def classify_tokens(prompt: str, labels: dict[str, str], tokens: list[str]):
+    tok_len = len(tokens)
+    example = "[" + (", ".join([f'"{tok}"' for tok in tokens])) + "]"
+    try:
+        response = await client.chat.completions.create(
+            model=OPENAI_MODEL,
+            timeout=30,
+            #temperature=0,
+            #presence_penalty=0,
+            reasoning_effort="low",
+            messages=[
+                {
+                    "role": "system",
+                    "content": (
+                        "Analyze the user provided sequence. Consider each string's semantic role "
+                        f"in the given sequence, then return a list of {tok_len} label strings. "
+                        f"Generate no more than {tok_len} labels. "
+                        "When typos or out-of-order words are provided, infer the intended meaning."
+                    ),
+                },
+                {
+                    "role": "system",
+                    "content": f"Labels: {labels}",
+                },
+                {
+                    "role": "user",
+                    "content": example,
+                },
+                {
+                    "role": "user",
+                    "content": (f"Replace each of the {tok_len} given strings with one of the following labels "
+                                f"that best describes {prompt}: {sorted(labels.keys())}"),
+                },
+            ],
+            response_format={
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "labels",
+                    "strict": True,
+                    "schema": {
+                        "type": "object",
+                        "properties": {
+                            "labels": {
+                                "type": "array",
+                                "description": f"List of {tok_len} labels, one for each string from the user's sequence.",
+                                "items": {
+                                    "type": "string",
+                                }
+                            }
+                        },
+                        "additionalProperties": False,
+                        "required": ["labels"]
+                    }
+                }
+            },
+        )
+    except Exception as e:
+        logger.error(f"openai call failed: {format_exc()}")
+        raise
+    try:
+        return json.loads(response.choices[0].message.content)["labels"]
+    except Exception as e:
+        logger.error(f"response: {response.choices[0].message} {format_exc()}")
+        raise
+
+
+async def classify_with_retry(prompt, labels, tokens, retry=10):
+    for i in range(retry):
+        try:
+            return await classify_tokens(prompt, labels, tokens)
+        except Exception as e:
+            logger.error(f"attempt {i} failed {tokens} {prompt} {format_exc()}")
+            await asyncio.sleep(i)
+
+
+async def generate_token_labels(case):
+    tokens = list(itertools.chain.from_iterable([s.strip("▁").split("▁") for s in sp.EncodeAsPieces(case)]))
+    (
+        adj,
+        adv,
+        det,
+        enc,
+        func,
+        misc,
+        ner1,
+        ner2,
+        noun,
+        pronoun,
+        punct,
+        verb,
+        wh,
+    ) = await asyncio.gather(
+        classify_with_retry(
+            f"its semantic role",
+            {
+                "JJ": "adjective",
+                "JJR": "comparative adjective",
+                "JJS": "superlative adjective",
+                "O": "out-of-scope",
+            },
+            tokens),
+        classify_with_retry(
+            f"its semantic role",
+            {
+                "RB": "adverb",
+                "RBR": "comparative adverb",
+                "RBS": "superlative adverb",
+                "O": "out-of-scope",
+            },
+            tokens),
+        classify_with_retry(
+            f"its semantic role",
+            {
+                "DT": "articles, demonstratives, and other determiners",
+                "EX": "existential 'there'",
+                "PDT": "predeterminer before a determiner to modify a noun phrase",
+                "O": "out-of-scope",
+            },
+            tokens),
+        classify_with_retry(
+            f"its sentence chunk classification",
+            {
+                "BRACKET": "in or contains bracket wrapped text",
+                "QUOTE": "in or contains quote wrapped text",
+                "TICK": "in or contains backtick wrapped text",
+                "O": "out-of-scope",
+            },
+            tokens),
+        classify_with_retry(
+            f"its semantic role",
+            {
+                "CC": "coordinating conjunction",
+                "IN": "preposition or subordinating conjunction",
+                "RP": "particle",
+                "TO": "to",
+                "UH": "interjection",
+                "O": "out-of-scope",
+            },
+            tokens),
+        classify_with_retry(
+            f"its semantic role",
+            {
+                "$": "currency",
+                "ADD": "address, URLs, usernames, or other non-lexical representations of places or entities",
+                "CD": "cardinal numbers",
+                "EMOJI": "emoji",
+                "TIME": "date or time",
+                "O": "out-of-scope",
+            },
+            tokens),
+        classify_with_retry(
+            f"its NER classification",
+            {
+                "B-GPE": "beginning of geopolitical entities",
+                "I-GPE": "inside of geopolitical entities",
+                "B-ORG": "beginning of organization",
+                "I-ORG": "inside of organization",
+                "B-PER": "beginning of person",
+                "I-PER": "inside of person",
+                "O": "out-of-scope",
+            },
+            tokens),
+        classify_with_retry(
+            f"its NER classification",
+            {
+                "B-EVENT": "beginning of event",
+                "I-EVENT": "inside of event",
+                "B-LOC": "beginning of location",
+                "I-LOC": "inside of location",
+                "O": "out-of-scope",
+            },
+            tokens),
+        classify_with_retry(
+            f"its semantic role",
+            {
+                "NN": "common noun singular",
+                "NNS": "common noun plural",
+                "NNP": "proper noun singular",
+                "NNPS": "proper noun plural",
+                "O": "out-of-scope",
+            },
+            tokens),
+        classify_with_retry(
+            f"its semantic role",
+            {
+                "POS": "possessive ending like the 's",
+                "PRP$": "possessive pronoun",
+                "PRP": "personal pronoun",
+                "O": "out-of-scope",
+            },
+            tokens),
+        classify_with_retry(
+            f"the punctuation classes it contains",
+            {
+                "COLON": "colon or semicolon",
+                "COMMA": "comma",
+                "EXCLAIM": "exclamation mark",
+                "HYPH": "dash or hyphen",
+                "LS": "list item marker",
+                "PERIOD": "period",
+                "QUESTION": "question mark",
+                "SEP": "any section separator",
+                "O": "out-of-scope",
+            },
+            tokens),
+        classify_with_retry(
+            f"its semantic role",
+            {
+                "MD": "modal verb",
+                "VB": "verb base form",
+                "VBD": "verb past tense",
+                "VBG": "present participle, gerund",
+                "VBN": "past participle",
+                "VBP": "non-3rd person singular present",
+                "VBZ": "3rd person singular present",
+                "O": "out-of-scope",
+            },
+            tokens),
+        classify_with_retry(
+            f"its semantic role",
+            {
+                "WDT": "Wh-determiner",
+                "WP$": "Wh-possessive pronoun",
+                "WP": "Wh-pronoun",
+                "WRB": "Wh-adverb",
+                "O": "out-of-scope",
+            },
+            tokens),
+        )
+    return {
+        "text": case,
+        "tokens": tokens,
+        "adj": adj,
+        "adv": adv,
+        "det": det,
+        "enc": enc,
+        "func": func,
+        "misc": misc,
+        "ner1": ner1,
+        "ner2": ner2,
+        "noun": noun,
+        "pronoun": pronoun,
+        "punct": punct,
+        "verb": verb,
+        "wh": wh,
+    }
+
+
+async def main(cases):
+    ds_dict = {
+        "text": [],
+        "tokens": [],
+        "adj": [],
+        "adv": [],
+        "det": [],
+        "enc": [],
+        "func": [],
+        "misc": [],
+        "ner1": [],
+        "ner2": [],
+        "noun": [],
+        "pronoun": [],
+        "punct": [],
+        "verb": [],
+        "wh": [],
+    }
+    drain_completed = False
+    max_concurrent_tasks = 15
+    tasks: list[Union[Task, None]] = []
+
+    async def checkpoint_task():
+        while not drain_completed:
+            # checkpoint
+            _ds = Dataset.from_dict(ds_dict)
+            _ds.save_to_disk("./custom_checkpoint_data")
+            logger.info(f"\n{_ds}")
+            await asyncio.sleep(600)
+    future_checkpoint_task_completion = asyncio.create_task(checkpoint_task())
+
+    async def drain_tasks():
+        while not drain_completed:
+            for idx, task in enumerate(tasks):
+                if task is None:
+                    continue
+                try:
+                    logger.info(f"attempting Example {idx}")
+                    example = await asyncio.wait_for(task, timeout=180)
+                    for col, labels in example.items():
+                        logger.info(f"  {col}: {labels}")
+                        ds_dict[col].append(example[col])
+                    tasks[idx] = None
+                except asyncio.exceptions.TimeoutError:
+                    logger.warning(f"attempt to wait_for Example {idx} timed out after 10 seconds.")
+                except Exception as e:
+                    logger.error(f"attempt to generate Example {idx} failed.\n{format_exc()}")
+                    tasks[idx] = None
+                    raise
+            await asyncio.sleep(1)
+    future_drain_completion = asyncio.create_task(drain_tasks())
+
+    for case in cases:
+        while len([t for t in tasks if t is not None]) >= max_concurrent_tasks:
+            await asyncio.sleep(1)
+        logger.info(f"scheduling case {case}")
+        tasks.append(asyncio.create_task(generate_token_labels(case)))
+
+    # Block until done
+    while len([t for t in tasks if t is not None]) > 0:
+        logger.info(f"waiting on {len([t for t in tasks if t is not None])} tasks")
+        await asyncio.sleep(1)
+
+    drain_completed = True
+    await future_drain_completion
+    await future_checkpoint_task_completion
+
+    ds = Dataset.from_dict(ds_dict)
+    ds.save_to_disk("./custom_final_data")
+    logger.info(f"\n{ds}")
+
+
+if __name__ == "__main__":
+    logging.config.dictConfig(default_logging_config)
+
+    all_examples = []
+
+    ud_en_ewt_ds = load_dataset("universal_dependencies", "en_ewt")
+    all_examples += ud_en_ewt_ds["train"]["text"]
+    all_examples += ud_en_ewt_ds["validation"]["text"]
+    all_examples += ud_en_ewt_ds["test"]["text"]
+
+    ud_en_gum_ds = load_dataset("universal_dependencies", "en_gum")
+    all_examples += ud_en_gum_ds["train"]["text"]
+    all_examples += ud_en_gum_ds["validation"]["text"]
+    all_examples += ud_en_gum_ds["test"]["text"]
+
+    ud_en_lines_ds = load_dataset("universal_dependencies", "en_lines")
+    all_examples += ud_en_lines_ds["train"]["text"]
+    all_examples += ud_en_lines_ds["validation"]["text"]
+    all_examples += ud_en_lines_ds["test"]["text"]
+
+    ud_en_partut_ds = load_dataset("universal_dependencies", "en_partut")
+    all_examples += ud_en_partut_ds["train"]["text"]
+    all_examples += ud_en_partut_ds["validation"]["text"]
+    all_examples += ud_en_partut_ds["test"]["text"]
+
+    ud_en_pronouns_ds = load_dataset("universal_dependencies", "en_pronouns")
+    all_examples += ud_en_pronouns_ds["test"]["text"]
+
+    ud_en_pud_ds = load_dataset("universal_dependencies", "en_pud")
+    all_examples += ud_en_pud_ds["test"]["text"]
+
+    logger.info(f"{len(all_examples)} UD examples")
+    random.shuffle(all_examples)
+    logger.info(f"{all_examples[:10]}")
+
+    all_examples += [
+        "Hello world!",
+        "127.0.0.1 is the localhost address.",
+        "1/2 is equivalent to 0.5 or 50%",
+        "John was running so fast, you can just tell he's a runner.",
+        "He excels at math and competed in the Math Olympiad",
+        "They're only $5!",
+        "Where is your sense of adventure?",
+        "I have only 3 cents.",
+        "Watson was on his way to 221B Baker Street when the robbery occurred.",
+        "That's uncopyrightable.",
+        "She's full of incomprehensibilities.",
+        "He's a total sesquipedalian.",
+        "you piece of SHIT!!",
+        "uh........... what..",
+        "Steph Curry is GOAT!!",
+        "[click here!](http://www.google.com)",
+        "Dude! The stock's value grew like 10x in a year!",
+        "Yea, I was at the DMV - God what a shit show!",
+        "Send an email to help@example.com",
+        "@goober, take your question to #corp-help-desk",
+        "Example 1 : Joe Shmoe has a big toe.",
+        "1. Steal under-pants. 2. ... 3. Profit!",
+        "I expect `len(word_list) == 3`",
+        "Home | Shop | Contact Us",
+        "This is me on cake <(^.^)>",
+        "and then he fell right on his face 😂",
+        "Putin is from Russia.",
+        "Zelenskyy is Ukrainian",
+        "In 2013, the Pentagon and other agencies officially acknowledged the existence of Area-51.",
+        "The Freedom of Information Act gives us the right to request access to records from any federal agency.",
+        "His motives here are totally sus",
+        "Yea, he finished by doing the Dab",
+        "Be back i'll",
+        "hes got a bon to pick",
+        "so then he says then he says, \"you'll regret this\" lol",
+    ]
+    asyncio.run(main(all_examples))
+

multi_task_classifier.py

@@ -0,0 +1,552 @@
+from datasets import DatasetDict, load_from_disk
+from sklearn.metrics import classification_report, precision_recall_fscore_support
+from transformers import (
+    DebertaV2Config,
+    DebertaV2Model,
+    DebertaV2PreTrainedModel,
+    DebertaV2TokenizerFast,
+    Trainer,
+    TrainingArguments,
+)
+import argparse
+import logging.config
+import numpy as np
+import torch
+import torch.nn as nn
+
+from ud_training_data_maker import get_uniq_training_labels, show_examples
+from utils import default_logging_config
+
+logger = logging.getLogger(__name__)
+
+arg_parser = argparse.ArgumentParser(description="Train multi-task model.")
+arg_parser.add_argument("-A", "--accumulation-steps", help="Gradient accumulation steps.",
+                        action="store", type=int, default=8)
+arg_parser.add_argument("--data-only", help='Show training data info and exit.',
+                        action="store_true", default=False)
+arg_parser.add_argument("--data-path", help="Load training dataset from specified path.",
+                        action="store", default="./training_data")
+arg_parser.add_argument("-E", "--train-epochs", help="Number of epochs to train for.",
+                        action="store", type=int, default=3)
+arg_parser.add_argument("-V", "--eval-batch-size", help="Per device eval batch size.",
+                        action="store", type=int, default=2)
+arg_parser.add_argument("--from-base", help="Load a base model.",
+                        action="store", default=None,
+                        choices=[
+                            "microsoft/deberta-v3-base",   # Requires --deberta-v3
+                            "microsoft/deberta-v3-large",  # Requires --deberta-v3
+                            # More?
+                        ])
+arg_parser.add_argument("-L", "--learning-rate", help="Learning rate.",
+                        action="store", type=float, default=5e-5)
+arg_parser.add_argument("--mini", help='Train model using small subset of examples for pipeline testing.',
+                        action="store_true", default=False)
+arg_parser.add_argument("--save-path", help="Save final model to specified path.",
+                        action="store", default="./final")
+arg_parser.add_argument("--show", help="Show examples: <split>/<col>/<label>/<count>",
+                        action="store", default=None)
+arg_parser.add_argument("--train", help='Train model using loaded examples.',
+                        action="store_true", default=False)
+arg_parser.add_argument("-T", "--train-batch-size", help="Per device train batch size.",
+                        action="store", type=int, default=2)
+args = arg_parser.parse_args()
+logging.config.dictConfig(default_logging_config)
+logger.info(f"Args {args}")
+
+
+
+# ------------------------------------------------------------------------------
+# Load dataset and show examples for manual inspection
+# ------------------------------------------------------------------------------
+
+loaded_dataset = load_from_disk(args.data_path)
+show_examples(loaded_dataset, args.show)
+
+# ------------------------------------------------------------------------------
+# Convert label analysis data into label sets for each head
+# ------------------------------------------------------------------------------
+
+ALL_LABELS = {col: list(vals) for col, vals in get_uniq_training_labels(loaded_dataset).items()}
+
+LABEL2ID = {
+    feat_name: {label: i for i, label in enumerate(ALL_LABELS[feat_name])}
+    for feat_name in ALL_LABELS
+}
+ID2LABEL = {
+    feat_name: {i: label for label, i in LABEL2ID[feat_name].items()}
+    for feat_name in LABEL2ID
+}
+
+# Each head's number of labels:
+NUM_LABELS_DICT = {k: len(v) for k, v in ALL_LABELS.items()}
+
+if args.data_only:
+    exit()
+
+# ------------------------------------------------------------------------------
+# Create a custom config that can store our multi-label info
+# ------------------------------------------------------------------------------
+
+class MultiHeadModelConfig(DebertaV2Config):
+    def __init__(self, label_maps=None, num_labels_dict=None, **kwargs):
+        super().__init__(**kwargs)
+        self.label_maps = label_maps or {}
+        self.num_labels_dict = num_labels_dict or {}
+
+    def to_dict(self):
+        output = super().to_dict()
+        output["label_maps"] = self.label_maps
+        output["num_labels_dict"] = self.num_labels_dict
+        return output
+
+# ------------------------------------------------------------------------------
+# Define a multi-head model
+# ------------------------------------------------------------------------------
+
+class MultiHeadModel(DebertaV2PreTrainedModel):
+    def __init__(self, config: MultiHeadModelConfig):
+        super().__init__(config)
+
+        self.deberta = DebertaV2Model(config)
+        self.classifiers = nn.ModuleDict()
+
+        hidden_size = config.hidden_size
+        for label_name, n_labels in config.num_labels_dict.items():
+            self.classifiers[label_name] = nn.Linear(hidden_size, n_labels)
+
+        # Initialize newly added weights
+        self.post_init()
+
+    def forward(
+            self,
+            input_ids=None,
+            attention_mask=None,
+            token_type_ids=None,
+            labels_dict=None,
+            **kwargs
+    ):
+        """
+        labels_dict: a dict of { label_name: (batch_size, seq_len) } with label ids.
+                     If provided, we compute and return the sum of CE losses.
+        """
+        outputs = self.deberta(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            **kwargs
+        )
+
+        sequence_output = outputs.last_hidden_state  # (batch_size, seq_len, hidden_size)
+
+        logits_dict = {}
+        for label_name, classifier in self.classifiers.items():
+            logits_dict[label_name] = classifier(sequence_output)
+
+        total_loss = None
+        loss_dict = {}
+        if labels_dict is not None:
+            # We'll sum the losses from each head
+            loss_fct = nn.CrossEntropyLoss()
+            total_loss = 0.0
+
+            for label_name, logits in logits_dict.items():
+                if label_name not in labels_dict:
+                    continue
+                label_ids = labels_dict[label_name]
+
+                # A typical approach for token classification:
+                # We ignore positions where label_ids == -100
+                active_loss = label_ids != -100  # shape (bs, seq_len)
+
+                # flatten everything
+                active_logits = logits.view(-1, logits.shape[-1])[active_loss.view(-1)]
+                active_labels = label_ids.view(-1)[active_loss.view(-1)]
+
+                loss = loss_fct(active_logits, active_labels)
+                loss_dict[label_name] = loss.item()
+                total_loss += loss
+
+        if labels_dict is not None:
+            # return (loss, predictions)
+            return total_loss, logits_dict
+        else:
+            # just return predictions
+            return logits_dict
+
+# ------------------------------------------------------------------------------
+# Tokenize with max_length=512, stride=128, and subword alignment
+# ------------------------------------------------------------------------------
+
+def tokenize_and_align_labels(examples):
+    """
+    For each example, the tokenizer may produce multiple overlapping
+    chunks if the tokens exceed 512 subwords. Each chunk will be
+    length=512, with a stride=128 for the next chunk.
+    We'll align labels so that subwords beyond the first in a token get -100.
+    """
+    # We rely on is_split_into_words=True because examples["tokens"] is a list of token strings.
+    tokenized_batch = tokenizer(
+        examples["tokens"],
+        is_split_into_words=True,
+        max_length=512,
+        stride=128,
+        truncation=True,
+        return_overflowing_tokens=True,
+        return_offsets_mapping=False,  # not mandatory for basic alignment
+        padding="max_length"
+    )
+
+    # The tokenizer returns "overflow_to_sample_mapping", telling us
+    # which original example index each chunk corresponds to.
+    # If the tokenizer didn't need to create overflows, the key might be missing
+    if "overflow_to_sample_mapping" not in tokenized_batch:
+        # No overflow => each input corresponds 1:1 with the original example
+        sample_map = [i for i in range(len(tokenized_batch["input_ids"]))]
+    else:
+        sample_map = tokenized_batch["overflow_to_sample_mapping"]
+
+    # We'll build lists for final outputs.
+    # For each chunk i, we produce:
+    #   "input_ids"[i], "attention_mask"[i], plus per-feature label IDs.
+    final_input_ids = []
+    final_attention_mask = []
+    final_labels_columns = {feat: [] for feat in ALL_LABELS}  # store one label-sequence per chunk
+
+    for i in range(len(tokenized_batch["input_ids"])):
+        # chunk i
+        chunk_input_ids = tokenized_batch["input_ids"][i]
+        chunk_attn_mask = tokenized_batch["attention_mask"][i]
+
+        original_index = sample_map[i]  # which example in the original batch
+        word_ids = tokenized_batch.word_ids(batch_index=i)
+
+        # We'll build label arrays for each feature
+        chunk_labels_dict = {}
+
+        for feat_name in ALL_LABELS:
+            # The UD token-level labels for the *original* example
+            token_labels = examples[feat_name][original_index]  # e.g. length T
+            chunk_label_ids = []
+
+            previous_word_id = None
+            for w_id in word_ids:
+                if w_id is None:
+                    # special token (CLS, SEP, padding)
+                    chunk_label_ids.append(-100)
+                else:
+                    # If it's the same word_id as before, it's a subword => label = -100
+                    if w_id == previous_word_id:
+                        chunk_label_ids.append(-100)
+                    else:
+                        # New token => use the actual label
+                        label_str = token_labels[w_id]
+                        label_id = LABEL2ID[feat_name][label_str]
+                        chunk_label_ids.append(label_id)
+                previous_word_id = w_id
+
+            chunk_labels_dict[feat_name] = chunk_label_ids
+
+        final_input_ids.append(chunk_input_ids)
+        final_attention_mask.append(chunk_attn_mask)
+        for feat_name in ALL_LABELS:
+            final_labels_columns[feat_name].append(chunk_labels_dict[feat_name])
+
+    # Return the new "flattened" set of chunks
+    # So the "map" call will expand each example → multiple chunk examples.
+    result = {
+        "input_ids": final_input_ids,
+        "attention_mask": final_attention_mask,
+    }
+    # We'll store each feature's label IDs in separate columns (e.g. labels_xpos, labels_deprel, etc.)
+    for feat_name in ALL_LABELS:
+        result[f"labels_{feat_name}"] = final_labels_columns[feat_name]
+
+    return result
+
+# ------------------------------------------------------------------------------
+# Trainer Setup
+# ------------------------------------------------------------------------------
+
+class MultiHeadTrainer(Trainer):
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        # 1) Gather all your per-feature labels from inputs
+        _labels_dict = {}
+        for feat_name in ALL_LABELS:
+            key = f"labels_{feat_name}"
+            if key in inputs:
+                _labels_dict[feat_name] = inputs[key]
+
+        # 2) Remove them so they don't get passed incorrectly to the model
+        for key in list(inputs.keys()):
+            if key.startswith("labels_"):
+                del inputs[key]
+
+        # 3) Call model(...) with _labels_dict
+        outputs = model(**inputs, labels_dict=_labels_dict)
+        # 'outputs' is (loss, logits_dict) in training/eval mode
+        loss, logits_dict = outputs
+
+        # Optional: if your special param is used upstream for some logic,
+        # you can handle it here or pass it along. For example:
+        if num_items_in_batch is not None:
+            # ... do something if needed ...
+            pass
+
+        if return_outputs:
+            # Return (loss, logits_dict) so Trainer sees logits_dict as predictions
+            return (loss, logits_dict)
+        else:
+            return loss
+
+    def prediction_step(self, model, inputs, prediction_loss_only=False, ignore_keys=None):
+        # 1) gather the "labels_xxx" columns
+        _labels_dict = {}
+        for feat_name in ALL_LABELS:
+            key = f"labels_{feat_name}"
+            if key in inputs:
+                _labels_dict[feat_name] = inputs[key]
+                del inputs[key]
+
+        # 2) forward pass without those keys
+        with torch.no_grad():
+            outputs = model(**inputs, labels_dict=_labels_dict)
+
+        loss, logits_dict = outputs  # you are returning (loss, dict-of-arrays)
+
+        if prediction_loss_only:
+            return (loss, None, None)
+
+        # The trainer expects a triple: (loss, predictions, labels)
+        #   - 'predictions' can be the dictionary
+        #   - 'labels' can be the dictionary of label IDs
+        return (loss, logits_dict, _labels_dict)
+
+
+def multi_head_classification_reports(logits_dict, labels_dict, id2label_dict):
+    """
+    For each head, generate a classification report (precision, recall, f1, etc. per class).
+    Return them as a dict: {head_name: "string report"}.
+    :param logits_dict: dict of {head_name: np.array(batch_size, seq_len, num_classes)}
+    :param labels_dict: dict of {head_name: np.array(batch_size, seq_len)}
+    :param id2label_dict: dict of {head_name: {id: label_str}}
+    :return: A dict of classification-report strings, one per head.
+    """
+    reports = {}
+
+    for head_name, logits in logits_dict.items():
+        if head_name not in labels_dict:
+            continue
+
+        predictions = np.argmax(logits, axis=-1)
+        valid_preds, valid_labels = [], []
+        for pred_seq, label_seq in zip(predictions, labels_dict[head_name]):
+            for p, lab in zip(pred_seq, label_seq):
+                if lab != -100:
+                    valid_preds.append(p)
+                    valid_labels.append(lab)
+
+        if len(valid_preds) == 0:
+            reports[head_name] = "No valid predictions."
+            continue
+
+        # Convert numeric IDs to string labels
+        valid_preds_str = [id2label_dict[head_name][p] for p in valid_preds]
+        valid_labels_str = [id2label_dict[head_name][l] for l in valid_labels]
+
+        # Generate the per-class classification report
+        report_str = classification_report(
+            valid_labels_str,
+            valid_preds_str,
+            zero_division=0
+        )
+        reports[head_name] = report_str
+
+    return reports
+
+
+def multi_head_compute_metrics(logits_dict, labels_dict):
+    """
+    For each head (e.g. xpos, deprel, Case, etc.), computes:
+      - Accuracy
+      - Precision (macro/micro)
+      - Recall (macro/micro)
+      - F1 (macro/micro)
+
+    :param logits_dict: dict of {head_name: np.array of shape (batch_size, seq_len, num_classes)}
+    :param labels_dict: dict of {head_name: np.array of shape (batch_size, seq_len)}
+    :return: A dict with aggregated metrics. Keys prefixed by head_name, e.g. "xpos_accuracy", "xpos_f1_macro", etc.
+    """
+    # We'll accumulate metrics in one big dictionary, keyed by "<head>_<metric>"
+    results = {}
+
+    for head_name, logits in logits_dict.items():
+        if head_name not in labels_dict:
+            # In case there's a mismatch or a head we didn't provide labels for
+            continue
+
+        # (batch_size, seq_len, num_classes)
+        predictions = np.argmax(logits, axis=-1)  # => (batch_size, seq_len)
+
+        # Flatten ignoring positions where label == -100
+        valid_preds, valid_labels = [], []
+        for pred_seq, label_seq in zip(predictions, labels_dict[head_name]):
+            for p, lab in zip(pred_seq, label_seq):
+                if lab != -100:
+                    valid_preds.append(p)
+                    valid_labels.append(lab)
+
+        valid_preds = np.array(valid_preds)
+        valid_labels = np.array(valid_labels)
+
+        if len(valid_preds) == 0:
+            # No valid data for this head—skip
+            continue
+
+        # Overall token-level accuracy
+        accuracy = (valid_preds == valid_labels).mean()
+
+        # Macro average => treat each class equally
+        precision_macro, recall_macro, f1_macro, _ = precision_recall_fscore_support(
+            valid_labels, valid_preds, average="macro", zero_division=0
+        )
+
+        # Micro average => aggregate across all classes
+        precision_micro, recall_micro, f1_micro, _ = precision_recall_fscore_support(
+            valid_labels, valid_preds, average="micro", zero_division=0
+        )
+
+        results[f"{head_name}_accuracy"] = accuracy
+        results[f"{head_name}_precision_macro"] = precision_macro
+        results[f"{head_name}_recall_macro"] = recall_macro
+        results[f"{head_name}_f1_macro"] = f1_macro
+        results[f"{head_name}_precision_micro"] = precision_micro
+        results[f"{head_name}_recall_micro"] = recall_micro
+        results[f"{head_name}_f1_micro"] = f1_micro
+
+    return results
+
+# ------------------------------------------------------------------------------
+# Instantiate model and tokenizer
+# ------------------------------------------------------------------------------
+
+if args.from_base:
+    model_name_or_path = args.from_base
+    multi_head_model = MultiHeadModel.from_pretrained(
+        model_name_or_path,
+        config=MultiHeadModelConfig.from_pretrained(
+            model_name_or_path,
+            num_labels_dict=NUM_LABELS_DICT,
+            label_maps=ALL_LABELS
+        )
+    )
+else:
+    model_name_or_path = args.save_path
+    # For evaluation, always load the saved checkpoint without overriding the config.
+    multi_head_model = MultiHeadModel.from_pretrained(model_name_or_path)
+    # EXTREMELY IMPORTANT!
+    # Override the label mapping based on the stored config to ensure consistency with training time ordering.
+    ALL_LABELS = multi_head_model.config.label_maps
+    LABEL2ID = {feat: {label: i for i, label in enumerate(ALL_LABELS[feat])} for feat in ALL_LABELS}
+    ID2LABEL = {feat: {i: label for label, i in LABEL2ID[feat].items()} for feat in LABEL2ID}
+logger.info(f"using {model_name_or_path}")
+
+# Check if GPU is usable
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif torch.backends.mps.is_available():  # For Apple Silicon MPS
+    device = torch.device("mps")
+else:
+    device = torch.device("cpu")
+logger.info(f"using {device}")
+multi_head_model.to(device)
+
+tokenizer = DebertaV2TokenizerFast.from_pretrained(
+    model_name_or_path,
+    add_prefix_space=True,
+)
+
+# ------------------------------------------------------------------------------
+# Shuffle, (optionally) sample, and tokenize final merged dataset
+# ------------------------------------------------------------------------------
+
+if args.mini:
+    loaded_dataset = DatasetDict({
+        "train": loaded_dataset["train"].shuffle(seed=42).select(range(1000)),
+        "validation": loaded_dataset["validation"].shuffle(seed=42).select(range(100)),
+        "test": loaded_dataset["test"].shuffle(seed=42).select(range(100)),
+    })
+
+# remove_columns => remove old "text", "tokens", etc. so we keep only model inputs
+tokenized_dataset = loaded_dataset.map(
+    tokenize_and_align_labels,
+    batched=True,
+    remove_columns=loaded_dataset["train"].column_names,
+)
+
+# ------------------------------------------------------------------------------
+# Train the model!
+# ------------------------------------------------------------------------------
+
+"""
+Current bests:
+
+deberta-v3-base:
+    num_train_epochs=3,
+    learning_rate=5e-5,
+    per_device_train_batch_size=2,
+    gradient_accumulation_steps=8,
+"""
+
+training_args = TrainingArguments(
+    # Evaluate less frequently or keep the same
+    eval_strategy="epoch",
+    num_train_epochs=args.train_epochs,
+    learning_rate=args.learning_rate,
+
+    output_dir="training_output",
+    overwrite_output_dir=True,
+    remove_unused_columns=False,  # important to keep the labels_xxx columns
+
+    logging_dir="training_logs",
+    logging_steps=100,
+
+    # Effective batch size = train_batch_size x gradient_accumulation_steps
+    per_device_train_batch_size=args.train_batch_size,
+    gradient_accumulation_steps=args.accumulation_steps,
+
+    per_device_eval_batch_size=args.eval_batch_size,
+)
+
+trainer = MultiHeadTrainer(
+    model=multi_head_model,
+    args=training_args,
+    train_dataset=tokenized_dataset["train"],
+    eval_dataset=tokenized_dataset["validation"],
+)
+
+if args.train:
+    trainer.train()
+    trainer.evaluate()
+    trainer.save_model(args.save_path)
+    tokenizer.save_pretrained(args.save_path)
+
+# ------------------------------------------------------------------------------
+# Evaluate the model!
+# ------------------------------------------------------------------------------
+
+pred_output = trainer.predict(tokenized_dataset["test"])
+pred_logits_dict = pred_output.predictions
+pred_labels_dict = pred_output.label_ids
+id2label_dict = ID2LABEL  # from earlier definitions
+
+# 1) Calculate metrics
+metrics = multi_head_compute_metrics(pred_logits_dict, pred_labels_dict)
+for k,v in metrics.items():
+    print(f"{k}: {v:.4f}")
+
+# 2) Print classification reports
+reports = multi_head_classification_reports(pred_logits_dict, pred_labels_dict, id2label_dict)
+for head_name, rstr in reports.items():
+    print(f"----- {head_name} classification report -----")
+    print(rstr)

requirements.txt

@@ -0,0 +1,16 @@
+accelerate>=0.26.0  # Required for using Trainer with PyTorch
+conllu
+dataset
+datasets
+numpy
+openai
+pytest
+pytest-cov
+pytest-xdist
+scikit-learn
+sentencepiece
+tensorflow
+tf-keras  # Required until transformers supports Keras 3
+tiktoken
+torch
+transformers

sp.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d2676ad813627497b95ce13c8ebe6b3313391c6df4b75909b5d6f68dcdde716b
+size 18104223

sp.vocab

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3a11823032d025ecd19a1e6bfef167b9a9ef6489d81eff726d4b399a20163ce
+size 18715604

sp_model_maker.py

@@ -0,0 +1,113 @@
+from datasets import load_dataset
+import argparse
+import logging.config
+import os
+import random
+import re
+import sentencepiece as spm
+
+from utils import default_logging_config
+
+logger = logging.getLogger(__name__)
+
+arg_parser = argparse.ArgumentParser(description="Train a sentencepiece tokenization model.")
+arg_parser.add_argument("--train", action="store_true", default=False,
+                        help="Train a sentencepiece tokenization model.")
+arg_parser.add_argument("--wikipedia", action="store_true", default=False,
+                        help="Use wikipedia dataset.")
+args = arg_parser.parse_args()
+logging.config.dictConfig(default_logging_config)
+
+input_sentence_size = 9_000_000
+max_line_char_len = 4192
+vocab_size = 900_000
+
+corpus_dir = "sp_data"
+corpus_file_prefix = f"{corpus_dir}/sp_corpus"
+model_file_prefix = "sp"
+uber_chunk_file = f"{corpus_dir}/wikipedia_uber_chunks.txt"
+white_space_pattern = re.compile(r"\s+")
+
+if args.wikipedia:
+    wikipedia_dataset_name = "20231101.en"
+    wikipedia_dataset = load_dataset("wikimedia/wikipedia", wikipedia_dataset_name)
+    total_page_cnt = len(wikipedia_dataset["train"])
+    logger.info(f"loaded {wikipedia_dataset_name} containing {total_page_cnt} pages")
+
+    max_processed_pages = total_page_cnt  # Change to single digits for spot checking / debugging
+    pages_processed_cnt = 0
+
+    corpus_file_part_idx = 0
+    current_corpus_file_char_len = 0
+    is_completed = False
+    iter_idx = 0
+    while not is_completed:  # Do till completed
+        with open(f"{corpus_file_prefix}_{corpus_file_part_idx}.txt", "a", encoding="utf-8") as f:
+            while iter_idx < (total_page_cnt - 1):
+                page = wikipedia_dataset["train"][iter_idx]
+                page_char_len = len(page["text"])  # Character len because bytes requires encoding
+                if page_char_len + current_corpus_file_char_len > 1_000_000_000:
+                    corpus_file_part_idx += 1  # New partition
+                    current_corpus_file_char_len = 0  # Reset tally
+                    break
+
+                page_chunk_cnt = 0
+                for page_chunk in page["text"].split("\n\n"):
+                    page_chunk_len = len(page_chunk)
+                    if not page_chunk or page_chunk[0] == " ":
+                        continue
+                    elif page_chunk_len > max_line_char_len:
+                        with open(uber_chunk_file, "a", encoding="utf-8") as uber_chunk_f:
+                            uber_chunk_f.write(page_chunk + "\n\n")
+                        continue
+
+                    page_chunk_lines = page_chunk.split("\n")
+                    for chunk_line in page_chunk_lines:
+                        if not chunk_line or chunk_line[0] == " ":
+                            continue
+                        elif len(white_space_pattern.split(chunk_line)) > 10:  # Require at least 10 naive tokens
+                            f.write(chunk_line + "\n")
+                            current_corpus_file_char_len += len(chunk_line)
+                    page_chunk_cnt += 1
+
+                iter_idx += 1
+                pages_processed_cnt += 1
+
+                if (pages_processed_cnt % 100) == 0:
+                    logger.info(f"processed {pages_processed_cnt}/{total_page_cnt} pages")
+                if pages_processed_cnt >= max_processed_pages:
+                    is_completed = True
+                    break
+            if not is_completed and iter_idx == (total_page_cnt - 1):
+                is_completed = True
+
+if args.train:
+    corpus_files = [f"{corpus_dir}/{f}" for f in os.listdir(corpus_dir) if f.startswith("sp_corpus")]
+    logger.info(f"corpus_files: {corpus_files}")
+
+    spm_training_args = [
+        f"--model_prefix={model_file_prefix}",
+        "--model_type=word",
+        "--shuffle_input_sentence=true",
+        #"--split_digits=true",
+        "--split_digits=false",
+        f"--input={','.join(random.sample(corpus_files, 15))}",
+        f"--input_sentence_size={input_sentence_size}",
+        f"--max_sentence_length={max_line_char_len}",
+        f"--vocab_size={vocab_size}",
+    ]
+    spm.SentencePieceTrainer.Train(" ".join(spm_training_args))
+
+# Now you can load the model and test it:
+sp = spm.SentencePieceProcessor()
+sp.LoadFromFile(f"{model_file_prefix}.model")
+
+print(sp.EncodeAsPieces("Hello world!"))
+print(sp.EncodeAsPieces("127.0.0.1 is the localhost address."))
+print(sp.EncodeAsPieces("1/2 is equivalent to 0.5 or 50%"))
+print(sp.EncodeAsPieces("John was running so fast, you can just tell he's a runner."))
+print(sp.EncodeAsPieces("He excels at math and competed in the Math Olympiad"))
+print(sp.EncodeAsPieces("Watson was on his way to 221B Baker Street when the robbery occurred."))
+print(sp.EncodeAsPieces("That's Uncopyrightable."))
+print(sp.EncodeAsPieces("She's full of incomprehensibilities."))
+print(sp.EncodeAsPieces("He's a total sesquipedalian."))

ud_dataset_maker.py

@@ -0,0 +1,427 @@
+from datasets import load_dataset, DatasetDict, concatenate_datasets
+from typing import Optional
+import argparse
+import ast
+import logging.config
+import random
+
+from utils.typos import generate_typo
+from utils import default_logging_config
+
+logger = logging.getLogger(__name__)
+
+allowed_xpos = [
+    "''",
+    '$',
+    ',',
+    '-LRB-',  # (
+    '-RRB-',  # )
+    '.',
+    ':',
+    'ADD',  # URLs, email addresses, or other “address” forms (like Twitter handles) that do not fit elsewhere.
+    'CC',
+    'CD',
+    'DT',
+    'EX',
+    'FW',
+    'HYPH',
+    'IN',
+    'JJ',
+    'JJR',
+    'JJS',
+    'LS',  # List item marker
+    'MD',
+    'NFP',  # “Non-Final Punctuation” for punctuation that doesn’t fit typical labels, in unexpected or stray positions
+    'NN',
+    'NNP',
+    'NNPS',
+    'NNS',
+    'PDT',
+    'POS',
+    'PRP$',
+    'PRP',
+    'RB',
+    'RBR',
+    'RBS',
+    'RP',
+    'SYM',
+    'TO',
+    'UH',
+    'VB',
+    'VBD',
+    'VBG',
+    'VBN',
+    'VBP',
+    'VBZ',
+    'WDT',
+    'WP$',
+    'WP',
+    'WRB',
+    '``',
+]
+
+allowed_deprel = [
+    'acl',
+    'acl:relcl',
+    'advcl',
+    'advmod',
+    'amod',
+    'appos',
+    'aux',
+    'aux:pass',
+    'case',
+    'cc',
+    'cc:preconj',
+    'ccomp',
+    'compound',
+    'compound:prt',
+    'conj',
+    'cop',
+    'csubj',
+    'csubj:pass',
+    'dep',
+    'det',
+    'det:predet',
+    'discourse',
+    'dislocated',
+    'expl',
+    'fixed',
+    'flat',
+    'flat:foreign',
+    'goeswith',
+    'iobj',
+    'list',
+    'mark',
+    'nmod',
+    'nmod:npmod',
+    'nmod:poss',
+    'nmod:tmod',
+    'nsubj',
+    'nsubj:pass',
+    'nummod',
+    'obj',
+    'obl',
+    'obl:npmod',
+    'obl:tmod',
+    'orphan',
+    'parataxis',
+    'punct',
+    'reparandum',
+    'root',
+    'vocative',
+    'xcomp',
+]
+
+target_feats = [
+    "Case", "Definite", "Degree", "Gender", "Mood", "NumType", "Number",
+    "Person", "Poss", "PronType", "Reflex", "Tense", "Typo", "VerbForm"
+]
+
+non_target_feats = {  # Found programmatically and added after analysis
+    "Abbr": [],
+    "Foreign": [],
+    "Polarity": [],
+    "Voice": [],
+}
+
+
+def add_target_feat_columns(exp):
+    """
+    Convert example["feats"] (list of feats) into separate columns
+    for each target_feat. Always return a dict with the same structure.
+    """
+    # example["feats"] is a list of length N (one per token)
+    feats_list = exp["feats"]
+
+    # Parse feats for each token
+    parsed_feats = [parse_morphological_feats(f, target_feats) for f in feats_list]
+
+    # Now add new columns for each target feat
+    for feat in target_feats:
+        exp[feat] = [pf[feat] for pf in parsed_feats]
+
+    return exp
+
+
+def get_uniq_training_labels(ds: DatasetDict, columns_to_exclude: set[str] = None):
+    columns_to_train_on = [k for k in ds["train"].features.keys() if k not in (
+        {"text", "tokens"} if columns_to_exclude is None else columns_to_exclude)]
+
+    # Create a dictionary of sets, keyed by each column name
+    label_counters = {col: dict() for col in columns_to_train_on}
+    unique_label_values = {col: set() for col in columns_to_train_on}
+
+    # Loop through each split and each example, and collect values
+    for split_name, dataset_split in ds.items():
+        for example in dataset_split:
+            # Each of these columns is a list (one entry per token),
+            # so we update our set with each token-level value
+            for col in columns_to_train_on:
+                unique_label_values[col].update(example[col])
+                for label_val in example[col]:
+                    if label_val not in label_counters[col]:
+                        label_counters[col][label_val] = 0  # Inits with 0
+                    label_counters[col][label_val] += 1
+
+    logger.info(f"Columns:")
+    for col in columns_to_train_on:
+        logger.info(f"  {col}:")
+        # Convert to a sorted list just to have a nice, stable ordering
+        vals = sorted(unique_label_values[col])
+        logger.info(f"    {len(vals)} labels: {[f'{v}:{label_counters[col][v]}' for v in vals]}")
+
+    return unique_label_values
+
+
+def introduce_typos(exp, typo_probability=0.03):
+    """
+    Randomly introduce typos in some % of tokens.
+    Update the `tokens` and the `Typo` columns in-place.
+    """
+    # new lists for mutated tokens and new Typo labels
+    mutated_tokens = []
+    mutated_typo_col = []
+
+    # Loop over each token
+    for token, old_typo_label in zip(exp["tokens"], exp["Typo"]):
+        # Decide whether to mutate this token
+        if random.random() < typo_probability:
+            mutated_token = generate_typo(token)
+            mutated_tokens.append(mutated_token)
+            mutated_typo_col.append("Yes")  # Mark as a "Yes" for the newly introduced typo
+        else:
+            mutated_tokens.append(token)
+            mutated_typo_col.append(old_typo_label)
+
+    exp["tokens"] = mutated_tokens
+    exp["Typo"] = mutated_typo_col
+    return exp
+
+
+def is_evenly_shaped(exp):
+    # All your target columns
+    feats = ["xpos", "deprel", *target_feats]
+    n_tokens = len(exp["tokens"])
+    for feat_name in feats:
+        if len(exp[feat_name]) != n_tokens:
+            return False
+    return True
+
+
+def is_valid_example(exp, dataset_name="ewt"):
+    """Return True if all xpos & deprel labels are in the common sets, else False."""
+    uniq_tokens = list(set(exp["tokens"]))
+    if len(uniq_tokens) == 1:
+        if uniq_tokens[0] == "_":
+            return False
+    for x in exp["xpos"]:
+        # If we hit an out-of-common-set xpos, we exclude this entire example
+        if x not in allowed_xpos:
+            # From time-to-time, we run into labels that are missing - either _ or None.
+            if x is None:
+                return False
+            elif x == "_":
+                return False
+            elif x == "-LSB-":  # [, en_gum only, not shared by other datasets
+                return False
+            elif x == "-RSB-":  # ], en_gum only, not shared by other datasets
+                return False
+            elif x == "AFX":  # “Affix” for bound morphemes or prefixes/suffixes that are split off from main tokens
+                return False
+            elif x == "GW":  # 'GW',  # "Gap Word", sometimes called “additional word” or “merged/gap word”).
+                return False
+            elif x == "XX":  # Unknown or “placeholder” words/tokens, 2 examples both word1/word2 with XX on the /
+                return False
+            logger.info(f"[{dataset_name}] Filtering example with: xpos={x}\n{exp['tokens']}\n{exp['xpos']}")
+            return False
+    for d in exp["deprel"]:
+        if d not in allowed_deprel:
+            if d is None:
+                return False
+            elif d == "_":
+                return False
+            logger.info(f"[{dataset_name}] Filtering example with: deprel={d}\n{exp['tokens']}\n{exp['deprel']}")
+            return False
+    return True
+
+
+def parse_morphological_feats(feats_in, targeted_feats):
+    """
+    Return a dict {feat_name: feat_value} for each target_feat.
+    If a feature is absent or doesn't apply, use "O".
+    If feats_in is a dict, read from it.
+    If feats_in is a string, parse it.
+    If feats_in is None/'_'/'' => no features => all "O".
+    """
+    # Default
+    out = {feat: "O" for feat in targeted_feats}
+
+    # Case A: feats_in is None or "_" or an empty string
+    if not feats_in or feats_in == "_" or feats_in == "None":
+        return out
+
+    pristine_feats_in = feats_in
+
+    # Case B: feats_in is a dict string: "{'Number': 'Sing', 'Person': '3'}"
+    if isinstance(feats_in, str):
+        feats_in = ast.literal_eval(feats_in)
+
+    # Case C: feats_in is a dictionary (some UD data does that)
+    if isinstance(feats_in, dict):
+        for k, v in feats_in.items():
+            if k in targeted_feats:
+                out[k] = v
+            else:
+                if k in non_target_feats:
+                    non_target_feats[k].append(v)
+                else:
+                    logger.info(f"Unhandled non-target feat '{k}={v}'")
+        return out
+
+    # Otherwise, unknown type
+    logger.warning(f"Unknown feats type {type(pristine_feats_in)} => {pristine_feats_in}")
+    return out
+
+
+def replace_bracket_label(exp):
+    label_map = {"(": "-LRB-", ")": "-RRB-"}
+    exp["xpos"] = [ label_map[tok] if tok in {"(", ")"} else tok for tok in exp["xpos"] ]
+    return exp
+
+
+def show_examples(ds: DatasetDict, show_expr: Optional[str]):
+    logger.info(f"Dataset:\n{ds}")
+    if not show_expr:
+        count_to_show = 2
+        examples_to_show = ds["train"][:count_to_show]
+    else:
+        args_show_tokens = show_expr.split("/")
+        split_to_show, col_to_show, label_to_show, count_to_show = args_show_tokens
+        count_to_show = int(count_to_show)
+        examples_to_show = ds[split_to_show].filter(
+            lambda exp: label_to_show in exp[col_to_show]).shuffle(seed=42)[:count_to_show]
+    for i in range(count_to_show):
+        logger.info(f"Example {i}:")
+        for feature in examples_to_show.keys():
+            logger.info(f"  {feature}: {examples_to_show[feature][i]}")
+
+
+def transform_and_filter_dataset(ud_dataset, dataset_name="ewt"):
+    """
+    ud_dataset is a DatasetDict with splits: 'train', 'validation', 'test' etc.
+    Return a new DatasetDict with the same splits but transformed/filtered.
+    """
+    new_splits = {}
+    for _split_name, _split_ds in ud_dataset.items():
+        if dataset_name == "pud":
+            _split_ds = _split_ds.map(replace_bracket_label)
+        filtered_split = _split_ds.filter(lambda ex: is_valid_example(ex, dataset_name=dataset_name))
+        transformed_split = filtered_split.map(
+            add_target_feat_columns,
+            batched=False
+        )
+        transformed_split = transformed_split.remove_columns(["deps", "feats", "head", "idx", "lemmas", "misc", "upos"])
+        new_splits[_split_name] = transformed_split.filter(is_evenly_shaped)
+    return DatasetDict(new_splits)
+
+
+if __name__ == "__main__":
+    arg_parser = argparse.ArgumentParser(description="Make training dataset.")
+    arg_parser.add_argument("--augment-typos", help='Augment final merged training data with typos.',
+                            action="store_true", default=False)
+    arg_parser.add_argument("--log-level", help='Log level.',
+                            action="store", default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"])
+    arg_parser.add_argument("--save", help='Save dataset to disk.',
+                            action="store_true", default=False)
+    arg_parser.add_argument("--save-path", help="Save final model to specified path.",
+                            action="store", default="./training_data")
+    arg_parser.add_argument("--show", help="Show examples: <split>/<col>/<label>/<count>",
+                            action="store", default=None)
+    args = arg_parser.parse_args()
+    logging.config.dictConfig(default_logging_config)
+
+    # Load UD Datasets: EWT, GUM, PUD
+    ud_en_ewt_ds = load_dataset("universal_dependencies", "en_ewt")
+    ud_en_gum_ds = load_dataset("universal_dependencies", "en_gum")
+    ud_en_pud_ds = load_dataset("universal_dependencies", "en_pud")
+
+    for loaded_ds_name, loaded_ds in {
+        "ud_en_ewt_ds": ud_en_ewt_ds,
+        "ud_en_gum_ds": ud_en_gum_ds,
+        "ud_en_pud_ds":  ud_en_pud_ds
+    }.items():
+        t_cnt = len(loaded_ds['test']) if 'test' in loaded_ds else 0
+        tr_cnt = len(loaded_ds['train']) if 'train' in loaded_ds else 0
+        v_cnt = len(loaded_ds['validation']) if 'train' in loaded_ds else 0
+        logger.info(f"Loaded {loaded_ds_name}: t:{t_cnt}, tr:{tr_cnt}, v:{v_cnt}")
+
+    # Apply transform + filtering to each split in each dataset
+    en_ewt_processed = transform_and_filter_dataset(ud_en_ewt_ds, "ewt")
+    en_gum_processed = transform_and_filter_dataset(ud_en_gum_ds, "gum")
+    en_pud_processed = transform_and_filter_dataset(ud_en_pud_ds, "pud")
+
+    def is_rare_case(exp):
+        if "ADD" in exp["xpos"]:
+            return True
+        if "LS" in exp["xpos"]:
+            return True
+        if "WP$" in exp["xpos"]:
+            return True
+        if "Cmp" in exp["Degree"]:
+            return True
+        if "Sup" in exp["Degree"]:
+            return True
+        if "Fem" in exp["Gender"]:
+            return True
+        if "Imp" in exp["Mood"]:
+            return True
+        if "Mult" in exp["NumType"]:
+            return True
+        if "Ord" in exp["NumType"]:
+            return True
+        if "1" in exp["Person"]:
+            return True
+        if "2" in exp["Person"]:
+            return True
+        if "Int" in exp["PronType"]:
+            return True
+        if "Rel" in exp["PronType"]:
+            return True
+        if "Yes" in exp["Reflex"]:
+            return True
+        if "Yes" in exp["Typo"]:
+            return True
+        if "Ger" in exp["VerbForm"]:
+            return True
+        return False
+
+    # Concatenate Datasets
+    final_dataset = DatasetDict()
+    final_dataset["test"] = concatenate_datasets(
+        [
+            en_ewt_processed["test"],
+            #en_gum_processed["test"].filter(is_rare_case),
+        ]
+    )
+
+    final_dataset["train"] = concatenate_datasets(
+        [
+            en_ewt_processed["train"],
+            #en_gum_processed["train"].filter(is_rare_case),
+            #en_pud_processed["test"].filter(is_rare_case),
+        ]
+    )
+    if args.augment_typos:
+        final_dataset["train"] = final_dataset["train"].map(introduce_typos, batched=False)
+
+    final_dataset["validation"] = concatenate_datasets(
+        [
+            en_ewt_processed["validation"],
+            #en_gum_processed["validation"].filter(is_rare_case),
+        ]
+    )
+    show_examples(final_dataset, args.show)
+    get_uniq_training_labels(final_dataset)
+    if args.save:
+        final_dataset.save_to_disk(args.save_path)
+

utils/__init__.py

@@ -0,0 +1,22 @@
+
+default_logging_config = {
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "default": {
+                "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+            },
+        },
+        "handlers": {
+            "console": {
+                "class": "logging.StreamHandler",
+                "formatter": "default",
+            },
+        },
+        "loggers": {
+            "": {
+                "level": "INFO",
+                "handlers": ["console"],
+            },
+        },
+    }

utils/typos.py

@@ -0,0 +1,417 @@
+import random
+import string
+
+common_misspelled_words = {
+    "absence": ["absense", "absentse", "abcense", "absance"],
+    "acceptable": ["acceptible"],
+    "accidentally": ["accidently", "ccidentaly"],
+    "accommodate": ["accomodate", "acommodate"],
+    "achieve": ["acheive"],
+    "acknowledge": ["acknowlege", "aknowledge"],
+    "acquaintance": ["acquaintence", "aquaintance"],
+    "acquire": ["aquire", "adquire"],
+    "acquit": ["aquit"],
+    "acreage": ["acrage", "acerage"],
+    "address": ["adress"],
+    "adultery": ["adultary"],
+    "advisable": ["adviseable", "advizable"],
+    "affect": ["effect"],
+    "aggression": ["agression"],
+    "aggressive": ["agressive"],
+    "allegiance": ["allegaince", "allegience", "alegiance"],
+    "almost": ["allmost"],
+    #"a lot": ["alot", "allot"]  # Not captured since "a lot" is two tokens.
+    "amateur": ["amatuer", "amature"],
+    "annually": ["anually", "annualy"],
+    "apparent": ["apparant", "aparent", "apparrent", "aparrent"],
+    "arctic": ["artic"],
+    "argument": ["arguement"],
+    "atheist": ["athiest", "athist"],
+    "awful": ["awfull", "aweful"],
+    "because": ["becuase", "becasue"],
+    "beautiful": ["beatiful"],
+    "becoming": ["becomeing"],
+    "beginning": ["begining"],
+    "believe": ["beleive"],
+    "bellwether": ["bellweather"],
+    "benefit": ["benifit"],
+    "buoy": ["bouy"],
+    "buoyant": ["bouyant"],
+    "business": ["buisness"],
+    "calendar": ["calender"],
+    "camouflage": ["camoflage", "camoflague"],
+    "capitol": ["capital"],
+    "Caribbean": ["Carribean"],  # More names?
+    "category": ["catagory"],
+    "caught": ["cauhgt", "caugt"],
+    "cemetery": ["cemetary", "cematery"],
+    "changeable": ["changable"],
+    "chief": ["cheif"],
+    "colleague": ["collaegue", "collegue"],
+    "column": ["colum"],
+    "coming": ["comming"],
+    "committed": ["commited", "comitted"],
+    "comparison": ["comparsion"],
+    "concede": ["conceed"],
+    "congratulate": ["congradulate"],
+    "conscientious": ["consciencious"],
+    "conscious": ["concious", "consious"],
+    "consensus": ["concensus"],
+    "controversy": ["contraversy"],
+    "coolly": ["cooly"],
+    "daiquiri": ["dacquiri", "daquiri"],
+    "deceive": ["decieve"],
+    "definite": ["definate", "definit"],
+    "definitely": ["definitly", "definately", "definatly", "defiantly"],
+    "desperate": ["desparate"],
+    "difference": ["diffrence"],
+    "dilemma": ["dilema"],
+    "disappoint": ["dissapoint"],
+    "disastrous": ["disasterous"],
+    "drunkenness": ["drunkeness"],
+    "dumbbell": ["dumbell"],
+    "embarrass": ["embarass"],
+    "equipment": ["equiptment"],
+    "exceed": ["excede"],
+    "exhilarate": ["exilerate"],
+    "existence": ["existance"],
+    "experience": ["experiance"],
+    "extreme": ["extreem"],
+    "fascinating": ["facinating"],
+    "fiery": ["firey"],
+    "fluorescent": ["flourescent"],
+    "foreign": ["foriegn"],
+    "forty": ["fourty"],
+    "friend": ["freind"],
+    "fulfil": ["fullfil", "fulfill"],
+    "gauge": ["guage"],
+    "grateful": ["gratefull", "greatful"],
+    "great": ["grate", "grat"],
+    "guarantee": ["garantee", "garentee", "garanty"],
+    "guidance": ["guidence"],
+    "harass": ["harrass"],
+    "height": ["heighth", "heigth"],
+    "hierarchy": ["heirarchy"],
+    # "hors d'oeuvres": ["hors derves", "ordeurves"]  # Not captured since "hors d'oeuvres" is two tokens.
+    "humorous": ["humerous"],
+    "hygiene": ["hygene", "hygine", "hiygeine", "higeine", "hygeine"],
+    "hypocrite": ["hipocrit"],
+    "ignorance": ["ignorence"],
+    "imitate": ["immitate"],
+    "immediately": ["imediately"],
+    "indict": ["indite"],
+    "independent": ["independant"],
+    "indispensable": ["indispensible"],
+    "inoculate": ["innoculate"],
+    "intelligence": ["inteligence", "intelligance"],
+    "jewelry": ["jewellery", "jewelery"],
+    "judgment": ["judgement"],
+    "kernel": ["kernal"],
+    "leisure": ["liesure"],
+    "liaison": ["liason"],
+    "library": ["libary", "liberry"],
+    "license": ["lisence", "licence"],
+    "lightning": ["lightening"],
+    "lose": ["loose"],
+    "maintenance": ["maintainance", "maintnance"],
+    "marshmallow": ["marshmellow"],
+    "medieval": ["medeval", "medevil", "mideval"],
+    "memento": ["momento"],
+    "millennium": ["millenium", "milennium"],
+    "miniature": ["miniture"],
+    "minuscule": ["miniscule"],
+    "mischievous": ["mischievious", "mischevous", "mischevious"],
+    "misspell": ["mispell", "misspel"],
+    "necessary": ["neccessary", "necessery"],
+    "niece": ["neice"],
+    "neighbour": ["nieghbor"],
+    "noticeable": ["noticable"],
+    "occasion": ["occassion"],
+    "occasionally": ["occasionaly", "occassionally"],
+    "occurrence": ["occurrance", "occurence"],
+    "occurred": ["occured"],
+    "omission": ["ommision", "omision"],
+    "original": ["orignal"],
+    "outrageous": ["outragous"],
+    "parliament": ["parliment"],
+    "pastime": ["passtime", "pasttime"],
+    "pedagogue": ["pedagoge"],
+    "perceive": ["percieve"],
+    "perseverance": ["perseverence"],
+    "personnel": ["personell", "personel"],
+    "plagiarize": ["plagerize"],
+    "playwright": ["playright", "playwrite"],
+    "possession": ["posession", "possesion"],
+    "potatoes": ["potatos"],
+    "precede": ["preceed"],
+    "presence": ["presance"],
+    "principle": ["principal"],
+    "privilege": ["privelege", "priviledge"],
+    "professor": ["professer"],
+    "protester": ["protestor"],
+    "promise": ["promiss"],
+    "pronunciation": ["pronounciatio"],
+    "proof": ["prufe"],
+    "prophecy": ["prophesy"],
+    "publicly": ["publically"],
+    "quarantine": ["quarentine"],
+    "queue": ["que"],
+    "questionnaire": ["questionaire", "questionnair"],
+    "readable": ["readible"],
+    "really": ["realy"],
+    "receive": ["recieve"],
+    "receipt": ["reciept"],
+    "recommend": ["recomend", "reccommend"],
+    "referred": ["refered"],
+    "reference": ["referance", "refrence"],
+    "relevant": ["relevent", "revelant"],
+    "religious": ["religous", "religius"],
+    "repetition": ["repitition"],
+    "restaurant": ["restarant", "restaraunt"],
+    "rhyme": ["rime"],
+    "rhythm": ["rythm", "rythem"],
+    "secretary": ["secratary", "secretery"],
+    "seize": ["sieze"],
+    "separate": ["seperate"],
+    "sergeant": ["sargent"],
+    "similar": ["similer"],
+    "skilful": ["skilfull", "skillful"],
+    "speech": ["speach", "speeche"],
+    "successful": ["succesful", "successfull", "sucessful"],
+    "supersede": ["supercede"],
+    "surprise": ["suprise", "surprize"],
+    "than": ["then"],
+    "their": ["there", "they're"],
+    "tomatoes": ["tomatos"],
+    "tomorrow": ["tommorow", "tommorrow"],
+    "Tucson": ["Tuscon"],
+    "twelfth": ["twelth"],
+    "tyranny": ["tyrany"],
+    "underrate": ["underate"],
+    "until": ["untill"],
+    "upholstery": ["upholstry"],
+    "usable": ["useable", "usible"],
+    "vacuum": ["vaccuum", "vaccum", "vacume"],
+    "vehicle": ["vehical"],
+    "vicious": ["visious"],
+    "what": ["wat"],
+    "weather": ["wether", "whether"],
+    "weird": ["wierd"],
+    "welfare": ["wellfare", "welfair"],
+    "whether": ["wether"],
+    "wilful": ["wilfull", "willful"],
+    "withhold": ["withold"],
+    "writing": ["writting", "writeing"],
+    "you're": ["your"],
+    "your": ["you're"],
+}
+
+
+def apostrophe_error(word: str) -> str:
+    """
+    Simulate common errors with apostrophes.
+
+    If the word contains an apostrophe:
+      - randomly remove it,
+      - shift it one position left (if possible),
+      - shift it one position right (if possible), or
+      - duplicate it.
+
+    If the word does not contain an apostrophe but ends with 's',
+    sometimes insert an apostrophe to mimic a mistaken possessive.
+    """
+    if "'" in word:
+        # Identify all apostrophe positions
+        indices = [i for i, ch in enumerate(word) if ch == "'"]
+        idx = random.choice(indices)
+        error_choice = random.choice(['remove', 'shift_left', 'shift_right', 'duplicate'])
+        if error_choice == 'remove':
+            return word[:idx] + word[idx + 1:]
+        elif error_choice == 'shift_left':
+            if idx > 0:
+                # Remove the apostrophe and insert it one position left.
+                return word[:idx - 1] + word[idx] + word[idx - 1] + word[idx + 1:]
+            else:
+                return word[:idx] + word[idx + 1:]
+        elif error_choice == 'shift_right':
+            if idx < len(word) - 1:
+                # Remove the apostrophe and insert it one position right.
+                return word[:idx] + word[idx + 1] + word[idx] + word[idx + 2:]
+            else:
+                return word[:idx] + word[idx + 1:]
+        elif error_choice == 'duplicate':
+            return word[:idx + 1] + "'" + word[idx + 1:]
+    else:
+        # For words without an apostrophe: if the word ends with 's', sometimes insert one.
+        if word.endswith("s") and random.random() < 0.5:
+            # Insert an apostrophe before the last letter.
+            return word[:-1] + "'" + word[-1]
+    return word
+
+
+def delete_random_letter(word: str) -> str:
+    """Simulate an omission error by deleting a random letter."""
+    if len(word) < 2:
+        return word
+    idx = random.randint(0, len(word) - 1)
+    return word[:idx] + word[idx + 1:]
+
+
+def duplicate_random_letter(word: str) -> str:
+    """Simulate an extra keypress by duplicating a letter at a random index."""
+    if not word:
+        return word
+    idx = random.randint(0, len(word) - 1)
+    return word[:idx + 1] + word[idx] + word[idx + 1:]
+
+
+def insert_random_letter(word: str) -> str:
+    """Simulate an insertion error by adding a random letter at a random position."""
+    idx = random.randint(0, len(word))
+    letter = random.choice(string.ascii_lowercase)
+    return word[:idx] + letter + word[idx:]
+
+
+def replace_with_adjacent_key(word: str) -> str:
+    """
+    Simulate a typing error by replacing a letter with one of its QWERTY neighbors.
+    Only letters with defined neighbors are considered.
+    """
+    # Define neighboring keys for a QWERTY keyboard (for lowercase letters)
+    qwerty_neighbors = {
+        'q': ['w', 'a'],
+        'w': ['q', 'e', 's'],
+        'e': ['w', 'r', 'd'],
+        'r': ['e', 't', 'f'],
+        't': ['r', 'y', 'g'],
+        'y': ['t', 'u', 'h'],
+        'u': ['y', 'i', 'j'],
+        'i': ['u', 'o', 'k'],
+        'o': ['i', 'p', 'l'],
+        'p': ['o'],
+        'a': ['q', 's', 'z'],
+        's': ['a', 'd', 'w', 'x'],
+        'd': ['s', 'f', 'e', 'c'],
+        'f': ['d', 'g', 'r', 'v'],
+        'g': ['f', 'h', 't', 'b'],
+        'h': ['g', 'j', 'y', 'n'],
+        'j': ['h', 'k', 'u', 'm'],
+        'k': ['j', 'l', 'i'],
+        'l': ['k', 'o'],
+        'z': ['a', 'x'],
+        'x': ['z', 'c', 's'],
+        'c': ['x', 'v', 'd'],
+        'v': ['c', 'b', 'f'],
+        'b': ['v', 'n', 'g'],
+        'n': ['b', 'm', 'h'],
+        'm': ['n', 'j']
+    }
+    # Find indices of characters that are letters with neighbors
+    valid_indices = [i for i, ch in enumerate(word) if ch.lower() in qwerty_neighbors]
+    if not valid_indices:
+        return word
+    idx = random.choice(valid_indices)
+    orig_char = word[idx]
+    lower_char = orig_char.lower()
+    replacement = random.choice(qwerty_neighbors[lower_char])
+    # Preserve original case
+    if orig_char.isupper():
+        replacement = replacement.upper()
+    return word[:idx] + replacement + word[idx + 1:]
+
+
+def swap_adjacent_letters(word: str) -> str:
+    """Simulate a transposition error by swapping two adjacent letters."""
+    if len(word) < 2:
+        return word
+    idx = random.randint(0, len(word) - 2)
+    word_list = list(word)
+    word_list[idx], word_list[idx + 1] = word_list[idx + 1], word_list[idx]
+    return ''.join(word_list)
+
+
+def switch_ie_ei(word: str) -> str:
+    """
+    Switch occurrences of 'ie' with 'ei' (or vice versa) to simulate
+    a common vowel pair error.
+    """
+    if 'ie' in word:
+        # Find all occurrences of 'ie'
+        indices = []
+        start = 0
+        while True:
+            idx = word.find('ie', start)
+            if idx == -1:
+                break
+            indices.append(idx)
+            start = idx + 1
+        if indices:
+            idx = random.choice(indices)
+            return word[:idx] + 'ei' + word[idx + 2:]
+    elif 'ei' in word:
+        indices = []
+        start = 0
+        while True:
+            idx = word.find('ei', start)
+            if idx == -1:
+                break
+            indices.append(idx)
+            start = idx + 1
+        if indices:
+            idx = random.choice(indices)
+            return word[:idx] + 'ie' + word[idx + 2:]
+    return word
+
+
+def generate_typo(word: str) -> str:
+    """
+    Given an input word, return a version of it with a common typo.
+    This function randomly selects one (or sometimes two) of the following error types:
+      - adjacent letter transposition
+      - deletion of a letter
+      - duplication of a letter
+      - insertion of a random letter
+      - replacement with a neighboring key (QWERTY)
+      - switching 'ie' and 'ei' sequences
+    While this method is by no means exhaustive, it reflects many of the typical errors documented.
+    """
+    if not word:
+        return word
+
+    if word in common_misspelled_words:
+        if random.random() < 0.5:  # 50% chance of selecting a common misspelling.
+            return random.choice(common_misspelled_words[word])
+
+    # List of available transformation functions
+    transformations = [
+        apostrophe_error,
+        delete_random_letter,
+        duplicate_random_letter,
+        insert_random_letter,
+        replace_with_adjacent_key,
+        swap_adjacent_letters,
+        switch_ie_ei
+    ]
+
+    # Randomly choose one transformation
+    transformation = random.choice(transformations)
+    result = transformation(word)
+
+    # Occasionally chain a second transformation (10% chance) for added variability
+    if random.random() < 0.1:
+        second_transformation = random.choice(transformations)
+        result = second_transformation(result)
+
+    return result
+
+
+# Example usage:
+if __name__ == '__main__':
+    test_words = [
+        "accommodate", "definitely", "receive", "mischievous", "calendar",
+        "equipment", "pronunciation", "consensus", "friend", "beautiful",
+        "doesn't", "books"
+    ]
+    for test_word in test_words:
+        typo = generate_typo(test_word)
+        print(f"Original: {test_word:15s} -> Typo: {typo}")