feat: emotions integration

goemotions_predict.py

@@ -0,0 +1,63 @@
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import numpy as np
+import torch
+
+from utils import get_torch_device
+
+
+class GoEmotionsPredictor:
+    def __init__(self, model_name_or_path: str, subfolder=None):
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path, subfolder=subfolder)
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            model_name_or_path, subfolder=subfolder)
+
+        self.label_names = getattr(self.model.config, "label_names", None)
+        self.per_label_thresh = getattr(self.model.config, "per_label_thresholds", None)
+        self.global_thresh = getattr(self.model.config, "best_global_threshold", 0.65)
+
+        self.device = get_torch_device()
+        self.model.to(self.device)
+        self.model.eval()
+
+    def predict(self, texts, use_per_label=True):
+        """
+        Args:
+          texts (list[str]): A list of raw text strings to classify.
+          use_per_label (bool): If True, apply per-label thresholds. If False, apply global threshold.
+        Returns:
+          A list of dicts, each with {"text": ..., "predicted_labels": [...]}
+        """
+        encodings = self.tokenizer(
+            texts,
+            truncation=True,
+            padding=True,
+            max_length=512,
+            return_tensors="pt"
+        )
+        # Move encodings to same device as the model
+        encodings = {k: v.to(self.device) for k, v in encodings.items()}
+
+        # 1) Run the model to get logits
+        with torch.no_grad():
+            outputs = self.model(**encodings)
+            logits = outputs.logits  # shape: (batch_size, num_labels)
+            probs = torch.sigmoid(logits).cpu().numpy()  # shape: (batch_size, num_labels)
+
+        # 2) Determine predictions by thresholding
+        if use_per_label:
+            # Use per-label thresholds
+            threshold_array = np.array(self.per_label_thresh)
+            preds = (probs >= threshold_array).astype(int)  # shape: (batch_size, num_labels)
+        else:
+            # Use global threshold
+            preds = (probs >= self.global_thresh).astype(int)
+
+        # 3) Convert integer predictions to label names
+        results = []
+        for i, text in enumerate(texts):
+            row_preds = preds[i]
+            predicted_labels = [self.label_names[j] for j, val in enumerate(row_preds) if val == 1]
+            results.append({"text": text, "emotions": predicted_labels})
+
+        return results

ud_dataset_maker.py

@@ -1,14 +1,21 @@
 from datasets import load_dataset, DatasetDict, concatenate_datasets
+from openai import OpenAI
+from traceback import format_exc
 import argparse
 import ast
+import json
 import logging.config
 import random
 
+from goemotions_predict import GoEmotionsPredictor
 from utils.typos import generate_typo
 from utils import default_logging_config, get_uniq_training_labels, show_examples
 
 logger = logging.getLogger(__name__)
 
+goemotions_predictor = GoEmotionsPredictor(
+    "veryfansome/deberta-goemotions", subfolder="pos_weight_best")
+
 allowed_xpos = [
     "''",
     '$',
@@ -111,11 +118,6 @@ allowed_deprel = [
     '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": [],
@@ -123,6 +125,68 @@ non_target_feats = {  # Found programmatically and added after analysis
     "Voice": [],
 }
 
+openai_classification_params = {
+    "model": "gpt-4o",
+    "temperature": 0.0,
+
+    #"model": "o3-mini",
+    #"reasoning_effort": "high",
+
+    "top_p": 1.0,
+    "presence_penalty": 0.0,
+    "frequency_penalty": 0.0,
+    "timeout": 30,
+}
+
+target_feats = [
+    "Case", "Definite", "Degree", "Gender", "Mood", "NumType", "Number",
+    "Person", "Poss", "PronType", "Reflex", "Tense", "Typo", "VerbForm"
+]
+
+word_lists_limiting_adjectives = [
+    "any",
+    "certain",
+    "each",
+    "every",
+    "other",
+    "some",
+
+    # Demonstrative adjectives / determiners
+    "that",
+    "these",
+    "this",
+    "those",
+]
+word_lists_difference_adjectives = [
+    "contrasting",
+    "different",
+    "disparate",
+    "dissimilar",
+    "distinct",
+    "divergent",
+    "diverse",
+    "heterogeneous",
+    "varied",
+    "various",
+]
+
+word_lists_similarity_adjectives = [
+    "alike",
+    "analogous",
+    "comparable",
+    "equal",
+    "equivalent",
+    "homogeneous",
+    "identical",
+    "interchangeable",
+    "same",
+    "similar",
+]
+
+word_lists_states_of_being_verbs = [
+    "am", "are", "be", "been", "being", "is", "was", "were",
+]
+
 
 def add_target_feat_columns(exp):
     """
@@ -142,6 +206,283 @@ def add_target_feat_columns(exp):
     return exp
 
 
+def extract_label_groups(exp, feat, target_labels=None):
+    """
+    For example, given a list of labels (e.g. ["O", "O", "NN", "NN", "O", "O", "NNS", "O"]),
+    this function will extract the index positions of the labels: NN, NNS, NNP, NNPS.
+
+    It returns a list of consecutive index groupings for those noun labels.
+    For example:
+        ["O", "O", "NN", "NN", "O", "O", "NNS", "O"]
+    would return:
+        [[2, 3], [6]]
+
+    Args:
+        exp: Example
+        feat: feature
+        target_labels (set of str): The set of tags to target.
+
+    Returns:
+        list of lists of int: A list where each sub-list contains consecutive indices
+                              of labels that match NN, NNS, NNP, NNPS.
+    """
+    groups = []
+    current_group = []
+
+    for idx, label in enumerate(exp[feat]):
+        if (label in target_labels) if target_labels is not None else label != "O":
+            # If current_group is empty or the current idx is consecutive (i.e., previous index + 1),
+            # append to current_group. Otherwise, start a new group.
+            if current_group and idx == current_group[-1] + 1:
+                current_group.append(idx)
+            else:
+                if current_group:
+                    groups.append(current_group)
+                current_group = [idx]
+        else:
+            if current_group:
+                groups.append(current_group)
+                current_group = []
+
+    # If there's an open group at the end, add it
+    if current_group:
+        groups.append(current_group)
+
+    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.
+""".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"{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"}]))
+    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
+label value, nothing else.
+""".replace("\n", "").strip()
+                                        },
+                                        {
+                                            "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"{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_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"""
+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
+{class_desc}. Use the B-{class_name_upper} label if the token begins a {class_name_upper} name entity and the
+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"]
+                                }
+                            }
+                        },
+                    )
+                    # 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()}")
+    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
+
+
 def introduce_typos(exp, typo_probability=0.03):
     """
     Randomly introduce typos in some % of tokens.
@@ -268,10 +609,34 @@ def transform_and_filter_dataset(ud_dataset, dataset_name="ewt"):
         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
         )
+        # 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"])
         new_splits[_split_name] = transformed_split.filter(is_evenly_shaped)
     return DatasetDict(new_splits)
@@ -312,55 +677,20 @@ if __name__ == "__main__":
     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),
-            en_pud_processed["test"],  #.filter(is_rare_case),
+            en_gum_processed["test"],
+            en_pud_processed["test"],
         ]
     )
 
     final_dataset["train"] = concatenate_datasets(
         [
             en_ewt_processed["train"],
-            en_gum_processed["train"],  #.filter(is_rare_case),
+            en_gum_processed["train"],
         ]
     )
     if args.augment_typos:
@@ -369,7 +699,7 @@ if __name__ == "__main__":
     final_dataset["validation"] = concatenate_datasets(
         [
             en_ewt_processed["validation"],
-            en_gum_processed["validation"],  #.filter(is_rare_case),
+            en_gum_processed["validation"],
         ]
     )
     show_examples(final_dataset, args.show)