wip: adj and adv features

ud_dataset_maker.py

@@ -1,4 +1,4 @@
-from datasets import load_dataset, DatasetDict, concatenate_datasets
+from datasets import load_dataset, load_from_disk, DatasetDict, concatenate_datasets
 from openai import OpenAI
 from traceback import format_exc
 import argparse
@@ -143,6 +143,40 @@ target_feats = [
     "Person", "Poss", "PronType", "Reflex", "Tense", "Typo", "VerbForm"
 ]
 
+word_lists_degree_adverbs = [
+    "almost",
+    "quite",
+    "rather",
+    "too",
+    "very",
+    "extremely",
+]
+
+word_lists_difference_adjectives = [
+    "contrasting",
+    "different",
+    "disparate",
+    "dissimilar",
+    "distinct",
+    "divergent",
+    "diverse",
+    "heterogeneous",
+    "varied",
+    "various",
+]
+
+word_lists_frequency_adverbs = [
+    "always",
+    "daily",
+    "monthly",
+    "often",
+    "rarely",
+    "seldom",
+    "sometimes",
+    "weekly",
+    "yearly",
+]
+
 word_lists_limiting_adjectives = [
     "any",
     "certain",
@@ -157,17 +191,9 @@ word_lists_limiting_adjectives = [
     "this",
     "those",
 ]
-word_lists_difference_adjectives = [
-    "contrasting",
-    "different",
-    "disparate",
-    "dissimilar",
-    "distinct",
-    "divergent",
-    "diverse",
-    "heterogeneous",
-    "varied",
-    "various",
+
+word_lists_negative_adverbs = [
+    "not",
 ]
 
 word_lists_similarity_adjectives = [
@@ -187,22 +213,77 @@ word_lists_states_of_being_verbs = [
     "am", "are", "be", "been", "being", "is", "was", "were",
 ]
 
+word_lists_time_adverbs = [
+    "already",
+    "soon",
+    "today",
+    "tomorrow",
+    "yesterday",
+]
+
+word_lists_uncertainty_adverbs = [
+    "maybe",
+    "perhaps",
+    "possibly",
+]
+
 
 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"]
+    if "feats" in exp:
+        # 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]
+        # 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
 
-    # Now add new columns for each target feat
-    for feat in target_feats:
-        exp[feat] = [pf[feat] for pf in parsed_feats]
 
+def convert_head_column(batch):
+    for feature_name, feature_attr in {
+        "AdjHead": ({"JJ", "JJR", "JJS"}, -4, 4),
+        "AdvHead": ({"RB", "RBR", "RBS"}, -3, 4),
+        "CdHead": ({"CD"}, -3, 3),
+        "ConjHead": ({"CC"}, -1, 4),
+        "DetHead": ({"DT", "PDT"}, -2, 4),
+        "InHead": ({"IN"}, -2, 5),
+        "ModalHead": ({"MD"}, -1, 3),
+        "NounHead": ({"NN", "NNS", "NNP", "NNPS"}, -5, 4),
+        "PronounHead": ({"PRP"}, -2, 3),
+        "ToHead": ({"TO"}, -1, 2),
+        "VerbHead": ({"VB", "VBD", "VBG", "VBN", "VBP", "VBZ"}, -5, 4),
+        "WhHead": ({"WDT", "WP", "WP$", "WRB"}, -2, 4),
+    }.items():
+        label_set, max_negative, max_positive = feature_attr
+        if feature_name not in batch:
+            batch[feature_name] = batch["head"].copy()
+            for head_idx, head_labels in enumerate(batch["head"]):
+                new_head_labels = []
+                for label_idx, label in enumerate(head_labels):
+                    if batch["xpos"][head_idx][label_idx] in label_set:
+                        new_label = int(label) - (label_idx + 1)
+                        if max_negative < new_label < max_positive:
+                            new_label = str(new_label)
+                        elif new_label > 0:
+                            new_label = f"{max_positive}+"
+                        else:
+                            new_label = f"{max_negative}+"
+                        new_head_labels.append(new_label)
+                    else:
+                        new_head_labels.append("O")
+                batch[feature_name][head_idx] = new_head_labels
+    return batch
+
+
+def convert_upos(exp, labels):
+    exp["pos"] = [labels[i] for i in exp.pop("upos")]
     return exp
 
 
@@ -251,152 +332,240 @@ def extract_label_groups(exp, feat, target_labels=None):
     return groups
 
 
-def introduce_emotion(exp):
-    exp["Emotion"] = ["X" for _ in exp["tokens"]]
-    labels = [l.upper() for l in goemotions_predictor.predict([exp["text"]], use_per_label=True)[0]["emotions"] if l != "neutral"]
-    labels.append("O")
-    labels_len = len(labels)
-    label_blob = ", ".join([(f"or {l}" if (labels_len > 1 and i == labels_len - 1) else l) for i, l in enumerate(labels)])
-    logger.info(f"label_blob: {label_blob}")
-    if label_blob != "O":
-        for capture_group in extract_label_groups(exp, "xpos", {
-            "JJ", "JJR", "JJS",
-            "NN", "NNS", "NNP", "NNPS",
-            "RB", "RBR", "RBS",
-            "VB", "VBD", "VBG", "VBN", "VBP", "VBZ",
-        }):
-            for token_idx in capture_group:
-                token = exp["tokens"][token_idx]
-                if token in word_lists_states_of_being_verbs:
-                    exp["Emotion"][token_idx] = "O"
-                else:
-                    with OpenAI() as client:
-                        while exp["Emotion"][token_idx] == "X":  # While not labeled
-                            try:
-                                completion = client.chat.completions.create(
-                                    messages=[
-                                        {
-                                            "role": "system",
-                                            "content": f"""
-Classify '{token}' at token index position {token_idx} by choosing the best fitting emotion label or O if out of scope.
-Pay close attention to semantic context but don't over-generalize if there is not enough context in the provided text. 
-Return only the label value, nothing else.
+def introduce_adj_type(exp):
+    if "AdjType" not in exp:
+        exp["AdjType"] = ["O" for _ in exp["tokens"]]
+        labels = ["Quantity", "Quality", "Size", "Age", "Shape", "Color", "Origin", "Material", "Purpose"]
+        labels_len = len(labels)
+        label_blob = ", ".join([(f"or {l}" if i == labels_len - 1 else l) for i, l in enumerate(labels)])
+        if "JJ" in exp["xpos"] or "JJR" in exp["xpos"] or "JJS" in exp["xpos"]:
+            for jj_group in extract_label_groups(exp, "xpos", {"JJ", "JJR", "JJS"}):
+                for jj_idx in jj_group:
+                    jj_token = exp["tokens"][jj_idx]
+                    if jj_token in word_lists_difference_adjectives:
+                        exp["AdjType"][jj_idx] = "Difference"
+                    elif jj_token in word_lists_limiting_adjectives:
+                        exp["AdjType"][jj_idx] = "Limit"
+                    elif jj_token in word_lists_similarity_adjectives:
+                        exp["AdjType"][jj_idx] = "Similarity"
+                    else:
+                        with OpenAI() as client:
+                            while exp["AdjType"][jj_idx] == "O":  # While not labeled
+                                try:
+                                    completion = client.chat.completions.create(
+                                        messages=[
+                                            {
+                                                "role": "system",
+                                                "content": f"""
+Classify '{jj_token}' at token index position {jj_idx} by choosing the best fitting adjective label. Return only the
+label value, nothing else.
 """.replace("\n", "").strip()
-                                        },
-                                        {
-                                            "role": "user",
-                                            "content": exp["text"]
-                                        },
-                                        {
-                                            "role": "user",
-                                            "content": str(exp["tokens"])
-                                        },
-                                        {
-                                            "role": "user",
-                                            "content": f"The word '{token}' at token index position {token_idx} above evokes {label_blob}?"
-                                        },
-                                    ],
-                                    **openai_classification_params,
-                                    response_format={
-                                        "type": "json_schema",
-                                        "json_schema": {
-                                            "name": "label",
-                                            "strict": True,
-                                            "schema": {
-                                                "type": "object",
-                                                "properties": {
-                                                    "label": {
-                                                        "type": "string",
-                                                        "enum": labels
-                                                    }
-                                                },
-                                                "additionalProperties": False,
-                                                "required": ["label"]
+                                            },
+                                            {
+                                                "role": "user",
+                                                "content": exp["text"]
+                                            },
+                                            {
+                                                "role": "user",
+                                                "content": str(exp["tokens"])
+                                            },
+                                            {
+                                                "role": "user",
+                                                "content": f"The adjective '{jj_token}' at token index position {jj_idx} above describes a {label_blob}?"
+                                            },
+                                        ],
+                                        **openai_classification_params,
+                                        response_format={
+                                            "type": "json_schema",
+                                            "json_schema": {
+                                                "name": "adjective",
+                                                "strict": True,
+                                                "schema": {
+                                                    "type": "object",
+                                                    "properties": {
+                                                        "label": {
+                                                            "type": "string",
+                                                            "enum": labels
+                                                        }
+                                                    },
+                                                    "additionalProperties": False,
+                                                    "required": ["label"]
+                                                }
                                             }
-                                        }
-                                    },
-                                )
-                                # Set so occasional hallucinations are retried
-                                new_label = json.loads(completion.choices[0].message.content)['label']
-                                logger.info(f"{token_idx}:{token} {new_label}")
-                                if new_label in labels:
-                                    exp["Emotion"][token_idx] = new_label
-                            except Exception as e:
-                                logger.error(f"failed to get label, trying again:\n{format_exc()}")
-        exp["Emotion"] = [("O" if l == "X" else l) for l in exp["Emotion"]]
-        logger.info("\n" + "\n".join([f"{k}\t{v}" for k, v in exp.items() if k in {"tokens", "Emotion"}]))
+                                        },
+                                    )
+                                    # Set so occasional hallucinations are retried
+                                    new_label = json.loads(completion.choices[0].message.content)['label']
+                                    logger.info(f"{jj_idx}:{jj_token} {new_label}")
+                                    if new_label in labels:
+                                        exp["AdjType"][jj_idx] = new_label
+                                except Exception as e:
+                                    logger.error(f"failed to get label, trying again:\n{format_exc()}")
+        logger.info("\n" + "\n".join([f"{k}\t{v}" for k, v in exp.items() if k in {"tokens", "AdjType"}]))
     return exp
 
 
-def introduce_adj_type(exp):
-    exp["AdjType"] = ["O" for _ in exp["tokens"]]
-    labels = ["Quantity", "Quality", "Size", "Age", "Shape", "Color", "Origin", "Material", "Purpose"]
-    labels_len = len(labels)
-    label_blob = ", ".join([(f"or {l}" if i == labels_len - 1 else l) for i, l in enumerate(labels)])
-    if "JJ" in exp["xpos"] or "JJR" in exp["xpos"] or "JJS" in exp["xpos"]:
-        for jj_group in extract_label_groups(exp, "xpos", {"JJ", "JJR", "JJS"}):
-            for jj_idx in jj_group:
-                jj_token = exp["tokens"][jj_idx]
-                if jj_token in word_lists_difference_adjectives:
-                    exp["AdjType"][jj_idx] = "Difference"
-                elif jj_token in word_lists_limiting_adjectives:
-                    exp["AdjType"][jj_idx] = "Limit"
-                elif jj_token in word_lists_similarity_adjectives:
-                    exp["AdjType"][jj_idx] = "Similarity"
-                else:
-                    with OpenAI() as client:
-                        while exp["AdjType"][jj_idx] == "O":  # While not labeled
-                            try:
-                                completion = client.chat.completions.create(
-                                    messages=[
-                                        {
-                                            "role": "system",
-                                            "content": f"""
-Classify '{jj_token}' at token index position {jj_idx} by choosing the best fitting adjective label. Return only the
+def introduce_adv_type(exp):
+    if "AdvType" not in exp:
+        exp["AdvType"] = ["O" for _ in exp["tokens"]]
+        labels = [
+            "Degree",
+            "Frequency",
+            "Manner",
+            "Negative",
+            "Place",
+            "Purpose",
+            "Time",
+            "Uncertainty",
+        ]
+        labels_len = len(labels)
+        label_blob = ", ".join([(f"or {l}" if i == labels_len - 1 else l) for i, l in enumerate(labels)])
+        if "RB" in exp["xpos"] or "RBR" in exp["xpos"] or "RBS" in exp["xpos"]:
+            for rb_group in extract_label_groups(exp, "xpos", {"RB", "RBR", "RBS"}):
+                for rb_idx in rb_group:
+                    rb_token = exp["tokens"][rb_idx]
+                    if rb_token in word_lists_degree_adverbs:
+                        exp["AdvType"][rb_idx] = "Degree"
+                    elif rb_token in word_lists_frequency_adverbs:
+                        exp["AdvType"][rb_idx] = "Frequency"
+                    elif rb_token in word_lists_negative_adverbs:
+                        exp["AdvType"][rb_idx] = "Negative"
+                    elif rb_token in word_lists_time_adverbs:
+                        exp["AdvType"][rb_idx] = "Time"
+                    elif rb_token in word_lists_uncertainty_adverbs:
+                        exp["AdvType"][rb_idx] = "Uncertainty"
+                    else:
+                        with OpenAI() as client:
+                            while exp["AdvType"][rb_idx] == "O":  # While not labeled
+                                try:
+                                    completion = client.chat.completions.create(
+                                        messages=[
+                                            {
+                                                "role": "system",
+                                                "content": f"""
+Classify '{rb_token}' at token index position {rb_idx} by choosing the best fitting adverb label. Return only the
 label value, nothing else.
 """.replace("\n", "").strip()
+                                            },
+                                            {
+                                                "role": "user",
+                                                "content": exp["text"]
+                                            },
+                                            {
+                                                "role": "user",
+                                                "content": str(exp["tokens"])
+                                            },
+                                            {
+                                                "role": "user",
+                                                "content": f"The adverb '{rb_token}' at token index position {rb_idx} above describes a {label_blob}?"
+                                            },
+                                        ],
+                                        **openai_classification_params,
+                                        response_format={
+                                            "type": "json_schema",
+                                            "json_schema": {
+                                                "name": "adverb",
+                                                "strict": True,
+                                                "schema": {
+                                                    "type": "object",
+                                                    "properties": {
+                                                        "label": {
+                                                            "type": "string",
+                                                            "enum": labels
+                                                        }
+                                                    },
+                                                    "additionalProperties": False,
+                                                    "required": ["label"]
+                                                }
+                                            }
                                         },
-                                        {
-                                            "role": "user",
-                                            "content": exp["text"]
-                                        },
-                                        {
-                                            "role": "user",
-                                            "content": str(exp["tokens"])
-                                        },
-                                        {
-                                            "role": "user",
-                                            "content": f"The adjective '{jj_token}' at token index position {jj_idx} above describes a {label_blob}?"
-                                        },
-                                    ],
-                                    **openai_classification_params,
-                                    response_format={
-                                        "type": "json_schema",
-                                        "json_schema": {
-                                            "name": "label",
-                                            "strict": True,
-                                            "schema": {
-                                                "type": "object",
-                                                "properties": {
-                                                    "label": {
-                                                        "type": "string",
-                                                        "enum": labels
-                                                    }
-                                                },
-                                                "additionalProperties": False,
-                                                "required": ["label"]
+                                    )
+                                    # Set so occasional hallucinations are retried
+                                    new_label = json.loads(completion.choices[0].message.content)['label']
+                                    logger.info(f"{rb_idx}:{rb_token} {new_label}")
+                                    if new_label in labels:
+                                        exp["AdvType"][rb_idx] = new_label
+                                except Exception as e:
+                                    logger.error(f"failed to get label, trying again:\n{format_exc()}")
+        logger.info("\n" + "\n".join([f"{k}\t{v}" for k, v in exp.items() if k in {"tokens", "AdvType"}]))
+    return exp
+
+
+def introduce_emotion(exp):
+    if "Emotion" not in exp:
+        exp["Emotion"] = ["X" for _ in exp["tokens"]]
+        labels = [l.upper() for l in goemotions_predictor.predict([exp["text"]], use_per_label=True)[0]["emotions"] if l != "neutral"]
+        labels.append("O")
+        labels_len = len(labels)
+        label_blob = ", ".join([(f"or {l}" if (labels_len > 1 and i == labels_len - 1) else l) for i, l in enumerate(labels)])
+        logger.info(f"label_blob: {label_blob}")
+        if label_blob != "O":
+            for capture_group in extract_label_groups(exp, "xpos", {
+                "JJ", "JJR", "JJS",
+                "NN", "NNS", "NNP", "NNPS",
+                "RB", "RBR", "RBS",
+                "VB", "VBD", "VBG", "VBN", "VBP", "VBZ",
+            }):
+                for token_idx in capture_group:
+                    token = exp["tokens"][token_idx]
+                    if token in word_lists_states_of_being_verbs:
+                        exp["Emotion"][token_idx] = "O"
+                    else:
+                        with OpenAI() as client:
+                            while exp["Emotion"][token_idx] == "X":  # While not labeled
+                                try:
+                                    completion = client.chat.completions.create(
+                                        messages=[
+                                            {
+                                                "role": "system",
+                                                "content": f"""
+Classify '{token}' at token index position {token_idx} by choosing the best fitting emotion label or O if out of scope.
+Pay close attention to semantic context but don't over-generalize if there is not enough context in the provided text. 
+Return only the label value, nothing else.
+""".replace("\n", "").strip()
+                                            },
+                                            {
+                                                "role": "user",
+                                                "content": exp["text"]
+                                            },
+                                            {
+                                                "role": "user",
+                                                "content": str(exp["tokens"])
+                                            },
+                                            {
+                                                "role": "user",
+                                                "content": f"The word '{token}' at token index position {token_idx} above evokes {label_blob}?"
+                                            },
+                                        ],
+                                        **openai_classification_params,
+                                        response_format={
+                                            "type": "json_schema",
+                                            "json_schema": {
+                                                "name": "label",
+                                                "strict": True,
+                                                "schema": {
+                                                    "type": "object",
+                                                    "properties": {
+                                                        "label": {
+                                                            "type": "string",
+                                                            "enum": labels
+                                                        }
+                                                    },
+                                                    "additionalProperties": False,
+                                                    "required": ["label"]
+                                                }
                                             }
-                                        }
-                                    },
-                                )
-                                # Set so occasional hallucinations are retried
-                                new_label = json.loads(completion.choices[0].message.content)['label']
-                                logger.info(f"{jj_idx}:{jj_token} {new_label}")
-                                if new_label in labels:
-                                    exp["AdjType"][jj_idx] = new_label
-                            except Exception as e:
-                                logger.error(f"failed to get label, trying again:\n{format_exc()}")
-    logger.info("\n" + "\n".join([f"{k}\t{v}" for k, v in exp.items() if k in {"tokens", "AdjType"}]))
+                                        },
+                                    )
+                                    # Set so occasional hallucinations are retried
+                                    new_label = json.loads(completion.choices[0].message.content)['label']
+                                    logger.info(f"{token_idx}:{token} {new_label}")
+                                    if new_label in labels:
+                                        exp["Emotion"][token_idx] = new_label
+                                except Exception as e:
+                                    logger.error(f"failed to get label, trying again:\n{format_exc()}")
+            logger.info("\n" + "\n".join([f"{k}\t{v}" for k, v in exp.items() if k in {"tokens", "Emotion"}]))
+    exp["Emotion"] = [("O" if l == "X" else l) for l in exp["Emotion"]]
     return exp
 
 
@@ -404,30 +573,32 @@ def introduce_ner_feature(exp, class_name: str, class_desc: str):
     class_name_capital = class_name.capitalize()
     class_name_upper = class_name.upper()
     class_feature_name = f"Ner{class_name_capital}"
-    exp[class_feature_name] = ["X" for _ in exp["tokens"]]
-
-    labels = [f"B-{class_name_upper}", f"I-{class_name_upper}", "O"]
-    labels_len = len(labels)
-    label_blob = ", ".join([(f"or {l}" if i == labels_len - 1 else l) for i, l in enumerate(labels)])
-    for capital_idx in [i for i, t in enumerate(exp["tokens"]) if len(t) > 0
-                                                                  and t[0].isupper()
-                                                                  and exp["xpos"][i] in {
-                                                                      "JJ", "JJR", "JJS",
-                                                                      "NN", "NNS", "NNP", "NNPS"
-                                                                  }]:
-        capital_token = exp["tokens"][capital_idx]
-        with OpenAI() as client:
-            while exp[class_feature_name][capital_idx] == "X":  # While not labeled
-                try:
-                    completion = client.chat.completions.create(
-                        messages=[
-                            {
-                                "role": "system",
-                                "content": "You are an expert in recognizing all kinds of names.",
-                            },
-                            {
-                                "role": "user",
-                                "content": f"""
+
+    if class_feature_name not in exp:
+        exp[class_feature_name] = ["X" for _ in exp["tokens"]]
+
+        labels = [f"B-{class_name_upper}", f"I-{class_name_upper}", "O"]
+        labels_len = len(labels)
+        label_blob = ", ".join([(f"or {l}" if i == labels_len - 1 else l) for i, l in enumerate(labels)])
+        for capital_idx in [i for i, t in enumerate(exp["tokens"]) if len(t) > 0
+                                                                      and t[0].isupper()
+                                                                      and exp["xpos"][i] in {
+                                                                          "JJ", "JJR", "JJS",
+                                                                          "NN", "NNS", "NNP", "NNPS"
+                                                                      }]:
+            capital_token = exp["tokens"][capital_idx]
+            with OpenAI() as client:
+                while exp[class_feature_name][capital_idx] == "X":  # While not labeled
+                    try:
+                        completion = client.chat.completions.create(
+                            messages=[
+                                {
+                                    "role": "system",
+                                    "content": "You are an expert in recognizing all kinds of names.",
+                                },
+                                {
+                                    "role": "user",
+                                    "content": f"""
 Classify '{capital_token}' at token index position {capital_idx} by choosing the best fitting BIO named entity label.
 Pay close attention to semantic context and neighboring tokens but don't over-generalize if there is not enough context
 in the provided text. Classify '{capital_token}' as a {class_name_upper} if it is being used as a part of a
@@ -435,51 +606,51 @@ in the provided text. Classify '{capital_token}' as a {class_name_upper} if it i
 I-{class_name_upper} label if '{capital_token}' continues a {class_name_upper} name entity. Return only the label
 value, nothing else.
 """.replace("\n", "").strip()
-                            },
-                            {
-                                "role": "user",
-                                "content": exp["text"]
-                            },
-                            {
-                                "role": "user",
-                                "content": str(exp["tokens"])
-                            },
-                            {
-                                "role": "user",
-                                "content": (f"The token '{capital_token}' at index position {capital_idx} above "
-                                            f"is used as a {label_blob} in the text?")
-                            },
-                        ],
-                        **openai_classification_params,
-                        response_format={
-                            "type": "json_schema",
-                            "json_schema": {
-                                "name": "label",
-                                "strict": True,
-                                "schema": {
-                                    "type": "object",
-                                    "properties": {
-                                        "label": {
-                                            "type": "string",
-                                            "enum": labels
-                                        }
-                                    },
-                                    "additionalProperties": False,
-                                    "required": ["label"]
+                                },
+                                {
+                                    "role": "user",
+                                    "content": exp["text"]
+                                },
+                                {
+                                    "role": "user",
+                                    "content": str(exp["tokens"])
+                                },
+                                {
+                                    "role": "user",
+                                    "content": (f"The token '{capital_token}' at index position {capital_idx} above "
+                                                f"is used as a {label_blob} in the text?")
+                                },
+                            ],
+                            **openai_classification_params,
+                            response_format={
+                                "type": "json_schema",
+                                "json_schema": {
+                                    "name": "label",
+                                    "strict": True,
+                                    "schema": {
+                                        "type": "object",
+                                        "properties": {
+                                            "label": {
+                                                "type": "string",
+                                                "enum": labels
+                                            }
+                                        },
+                                        "additionalProperties": False,
+                                        "required": ["label"]
+                                    }
                                 }
-                            }
-                        },
-                    )
-                    # Set if valid label so occasional hallucinations are retried
-                    new_label = json.loads(completion.choices[0].message.content)['label']
-                    logger.info(f"{capital_idx}:{capital_token} {new_label}")
-                    if new_label in labels:
-                        exp[class_feature_name][capital_idx] = new_label
-                except Exception as e:
-                    logger.error(f"failed to get {class_feature_name} label for {capital_token} at idx {capital_idx} "
-                                 f"in \"{exp['text']}\", trying again:\n{format_exc()}")
+                            },
+                        )
+                        # Set if valid label so occasional hallucinations are retried
+                        new_label = json.loads(completion.choices[0].message.content)['label']
+                        logger.info(f"{capital_idx}:{capital_token} {new_label}")
+                        if new_label in labels:
+                            exp[class_feature_name][capital_idx] = new_label
+                    except Exception as e:
+                        logger.error(f"failed to get {class_feature_name} label for {capital_token} at idx {capital_idx} "
+                                     f"in \"{exp['text']}\", trying again:\n{format_exc()}")
+            logger.info("\n" + "\n".join([f"{k}\t{v}" for k, v in exp.items() if k in {"tokens", class_feature_name}]))
     exp[class_feature_name] = [("O" if l == "X" else l) for l in exp[class_feature_name]]
-    logger.info("\n" + "\n".join([f"{k}\t{v}" for k, v in exp.items() if k in {"tokens", class_feature_name}]))
     return exp
 
 
@@ -610,34 +781,38 @@ def transform_and_filter_dataset(ud_dataset, dataset_name="ewt"):
             _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(
+        transformed_split = filtered_split.map(lambda exp: convert_upos(exp, _split_ds.features["upos"].feature.names),
+                                               batched=False)
+        transformed_split = transformed_split.map(
             add_target_feat_columns,
             batched=False
         )
+        transformed_split = transformed_split.map(convert_head_column, batched=True, batch_size=1000)
         # TODO:
         #   - Get emotion classes and label adj and adv tokens based on classified emotions. This connects descriptions,
         #     with the kind of attribute, with the emotions evoked.
         #   - checkpoints after each phase to avoid costly re-dos
-        transformed_split = transformed_split.map(introduce_emotion, batched=False)
-        transformed_split = transformed_split.map(introduce_adj_type, batched=False)
-        transformed_split = transformed_split.map(
-            lambda exp: introduce_ner_feature(
-                exp, "location",
-                "location's name"),
-            batched=False)
-        transformed_split = transformed_split.map(
-            lambda exp: introduce_ner_feature(
-                exp, "organization",
-                "organization's name"),
-            batched=False)
-        transformed_split = transformed_split.map(
-            lambda exp: introduce_ner_feature(
-                exp, "person",
-                "person's name"),
-            batched=False)
-
-        new_splits[_split_name] = transformed_split
-        transformed_split = transformed_split.remove_columns(["deps", "feats", "head", "idx", "lemmas", "misc", "upos"])
+        #transformed_split = transformed_split.map(introduce_emotion, batched=False)
+        #transformed_split = transformed_split.map(introduce_adj_type, batched=False)
+        #transformed_split = transformed_split.map(
+        #    lambda exp: introduce_ner_feature(
+        #        exp, "location",
+        #        "location's name"),
+        #    batched=False)
+        #transformed_split = transformed_split.map(
+        #    lambda exp: introduce_ner_feature(
+        #        exp, "organization",
+        #        "organization's name"),
+        #    batched=False)
+        #transformed_split = transformed_split.map(
+        #    lambda exp: introduce_ner_feature(
+        #        exp, "person",
+        #        "person's name"),
+        #    batched=False)
+
+        for col_name in {"deps", "feats", "head", "idx", "lemmas", "misc"}:
+            if col_name in transformed_split.features:
+                transformed_split = transformed_split.remove_columns([col_name])
         new_splits[_split_name] = transformed_split.filter(is_evenly_shaped)
     return DatasetDict(new_splits)
 
@@ -646,62 +821,68 @@ 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("--load-path", help="Load dataset from specified path.",
+                            action="store", default=None)
     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.",
+    arg_parser.add_argument("--save-path", help="Save final dataset to specified path.",
                             action="store", default="./ud_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")
+    if args.load_path is None:
+        # 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")
+
+        # Concatenate Datasets
+        final_dataset = DatasetDict()
+        final_dataset["test"] = concatenate_datasets(
+            [
+                en_ewt_processed["test"],
+                en_gum_processed["test"],
+                en_pud_processed["test"],
+            ]
+        )
 
-    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")
-
-    # Concatenate Datasets
-    final_dataset = DatasetDict()
-    final_dataset["test"] = concatenate_datasets(
-        [
-            en_ewt_processed["test"],
-            en_gum_processed["test"],
-            en_pud_processed["test"],
-        ]
-    )
+        final_dataset["train"] = concatenate_datasets(
+            [
+                en_ewt_processed["train"],
+                en_gum_processed["train"],
+            ]
+        )
+        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"],
+            ]
+        )
+    else:
+        final_dataset = transform_and_filter_dataset(load_from_disk(args.load_path))
 
-    final_dataset["train"] = concatenate_datasets(
-        [
-            en_ewt_processed["train"],
-            en_gum_processed["train"],
-        ]
-    )
-    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"],
-        ]
-    )
     show_examples(final_dataset, args.show)
     get_uniq_training_labels(final_dataset)
     if args.save: