baseline_model.py

import csv
import os
from typing import List, Dict, Tuple

import numpy as np
from sklearn.feature_extraction import DictVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.model_selection import GroupKFold
from sklearn.metrics import classification_report

from helper_functions.baseline_model_helper_functions import read_conllu_folder, collapse_label
from paths import gold_standard_folder_name, baseline_output_folder_name


def char_ngrams(word: str, n: int = 3) -> List[str]:
    w = f"^{word}$"
    res = []
    for i in range(len(w) - n + 1):
        res.append(w[i:i+n])
    return res


def extract_token_features(sentence: List[Dict], i: int, prev_n: int = 5) -> Dict[str, int]:
    tok = sentence[i]['token']
    lower = tok.lower()

    # basic token-level features
    token_features = {f"tok_lower={lower}": 1, f"tok_is_title={tok.istitle()}": 1,  f"tok_is_upper={tok.isupper()}": 1,
                      f"tok_has_digit={any(c.isdigit() for c in tok)}": 1, f"pos_index={i}": 1}

    # features from previous tokens in the sentence (left context)
    for p in range(1, prev_n + 1):
        idx = i - p

        if idx >= 0:
            prev_tok = sentence[idx]['token']
            prev_lower = prev_tok.lower()

            token_features[f"prev{p}_lower={prev_lower}"] = 1
            token_features[f"prev{p}_is_upper={prev_tok.isupper()}"] = 1

        else:
            token_features[f"prev{p}_BOS"] = 1

    # short prefixes / suffixes to capture subword patterns
    if len(lower) >= 1:
        token_features[f"pref1={lower[:1]}"] = 1
    if len(lower) >= 2:
        token_features[f"pref2={lower[:2]}"] = 1
    if len(lower) >= 3:
        token_features[f"pref3={lower[:3]}"] = 1
    if len(lower) >= 1:
        token_features[f"suf1={lower[-1:]}"] = 1
    if len(lower) >= 2:
        token_features[f"suf2={lower[-2:]}"] = 1
    if len(lower) >= 3:
        token_features[f"suf3={lower[-3:]}"] = 1

    # character trigrams as additional robust features.
    for ng in char_ngrams(lower, 3):
        token_features[f"cng3={ng}"] = 1

    return token_features


def metrics_per_class(y_test, y_pred, variant, fold):
    # metrics per class
    rows = []
    for cls, metrics in classification_report(y_test, y_pred, zero_division=0, output_dict=True).items():
        if cls == "accuracy":
            rows.append({
                "fold": fold, "variant": variant, "class": "accuracy",
                "precision": metrics, "recall": metrics, "f1": metrics,
                "support": float(len(y_test)) # to be the same type as others
            })
        elif isinstance(metrics, dict):
            rows.append({
                "fold": fold, "variant": variant, "class": cls,
                "precision": metrics.get("precision", 0.0),
                "recall": metrics.get("recall", 0.0),
                "f1": metrics.get("f1-score", 0.0),
                "support": metrics.get("support", 0)
            })
    return rows


def train_and_evaluate_fold(x_train, x_test, y_train, y_test, variant: str, fold: int) -> Tuple[List[Dict], List[str], List[str]]:
    clf = MultinomialNB()

    clf.fit(x_train, y_train)
    predictions = clf.predict(x_test)

    return metrics_per_class(y_test, predictions, variant, fold), list(y_test), list(predictions)


def build_dataset(sentences: List[List[Dict]], prev_n: int = 2):
    X_dicts, Y, groups = [], [], []

    for s_idx, sent in enumerate(sentences):
        for i in range(len(sent)):
            feats = extract_token_features(sent, i, prev_n=prev_n)

            X_dicts.append(feats)
            Y.append(sent[i]["ner"])
            groups.append(s_idx)

    return np.array(X_dicts, dtype=object), np.array(Y, dtype=object), np.array(groups, dtype=int)


def save_overall_report(y_true, y_pred, variant: str, output_dir: str, sentences=None, groups=None):
    os.makedirs(output_dir, exist_ok=True)
    txt_path = os.path.join(output_dir, f"token_level_report_{variant}.txt")

    with open(txt_path, "w", encoding="utf-8") as f:

        if sentences is not None and groups is not None:
            sent_to_domain = {i: s[0].get("domain", "UNKNOWN") for i, s in enumerate(sentences)}

            domain_data = {}
            for tok_idx, sent_idx in enumerate(groups):
                domain = sent_to_domain.get(sent_idx, "UNKNOWN")
                domain_data.setdefault(domain, {"y_true": [], "y_pred": []})
                domain_data[domain]["y_true"].append(y_true[tok_idx])
                domain_data[domain]["y_pred"].append(y_pred[tok_idx])

            for domain, vals in domain_data.items():
                report = classification_report(vals["y_true"], vals["y_pred"], zero_division=0, digits=2)
                header = f"{domain.upper()} Domain:\n\n"
                print(header + report + "\n")
                f.write(header)
                f.write(report)
                f.write("\n\n")

        overall_report = classification_report(y_true, y_pred, zero_division=0, digits=2)
        print(f"ALL Domains Together:\n\n{overall_report}\n")
        f.write("ALL Domains Together:\n\n")
        f.write(overall_report)
        f.write("\n")


def run_cross_validation(sentences: List[List[Dict]], output_dir: str, n_splits: int = 10, prev_n: int = 5):
    os.makedirs(output_dir, exist_ok=True)

    # build token-level dataset with sentence indices used as groups (for GroupKFold)
    x_array, y_array, groups_array = build_dataset(sentences, prev_n=prev_n)

    gkf = GroupKFold(n_splits=n_splits)

    token_rows = [] # per-fold per-class metric rows

    all_y_true_bio, all_predictions_bio, all_y_true_coll, all_predictions_coll = [], [], [], []

    for fold, (train_idx, test_idx) in enumerate(list(gkf.split(x_array, y_array, groups_array))):

        # fitting vectorizer only on train
        dv = DictVectorizer(sparse=True)
        X_train = dv.fit_transform(x_array[train_idx])
        X_test = dv.transform(x_array[test_idx])


        # BIO (whole labels)
        y_train_bio = y_array[train_idx]
        y_test_bio = y_array[test_idx]

        bio_rows, y_true_bio, y_pred_bio = train_and_evaluate_fold(X_train, X_test, y_train_bio, y_test_bio, variant="bio", fold=fold)

        token_rows.extend(bio_rows)

        all_y_true_bio.extend(y_true_bio)
        all_predictions_bio.extend(y_pred_bio)


        # collapsed (just simplified labels)
        y_train_coll = np.array([collapse_label(l) for l in y_train_bio])
        y_test_coll = np.array([collapse_label(l) for l in y_test_bio])

        coll_rows, y_true_coll, y_pred_coll = train_and_evaluate_fold(X_train, X_test, y_train_coll, y_test_coll, variant="collapsed", fold=fold)

        token_rows.extend(coll_rows)

        all_y_true_coll.extend(y_true_coll)
        all_predictions_coll.extend(y_pred_coll)

    # saving per-fold per class CSV - both variants
    with open(os.path.join(output_dir, 'token_level_results_per_fold.csv'), mode="w", newline="", encoding="utf-8") as f:
        if token_rows:
            writer = csv.DictWriter(f, fieldnames=list(token_rows[0].keys()))
            writer.writeheader()
            writer.writerows(token_rows)

    # overall aggregated reports
    if all_y_true_bio:
        save_overall_report(all_y_true_bio, all_predictions_bio, 'bio', output_dir, sentences, groups_array)

    if all_y_true_coll:
        save_overall_report(all_y_true_coll, all_predictions_coll, 'collapsed', output_dir, sentences, groups_array)


def main():
    # Baseline NER: 10-fold cross-validation over the gold standard dataset with a simple NB + feature-based model
    n_splits = 10
    prev_n = 2
    sentences = read_conllu_folder(gold_standard_folder_name)
    if not sentences:
        print("No sentences read. Couldn't perform a cross validation")
        return
    print(f'Read {len(sentences)} sentences, example sentence length: {len(sentences[0])}')
    run_cross_validation(sentences, baseline_output_folder_name, n_splits=n_splits, prev_n=prev_n)


if __name__ == '__main__':
    main()