add split creation script and fix typos

Author: whoisjones

README.md

@@ -1,20 +1,20 @@
-# Label Shift Estimation for Named Entity Recognition using Familarity
+# Label Shift Estimation for Named Entity Recognition using Familiarity
 
 **Our paper got accepted to NAACL 2025 🎉 See our [paper](https://arxiv.org/abs/2412.10121) and find the datasets on the [huggingface hub]()!**
 
-This repository computes the label shift for zero-shot NER settings using the Familarity metric. The metric uses semantic similarity between the sets of label seen during training and used for evaluation to indicate how "familiar" the trained model will be with the evaluation labels.
+This repository computes the label shift for zero-shot NER settings using the Familiarity metric. The metric uses semantic similarity between the sets of label seen during training and used for evaluation to indicate how "familiar" the trained model will be with the evaluation labels.
 
 ## Installation
 ```python
-conda create -n familarity python=3.11
-conda activate familarity
+conda create -n familiarity python=3.11
+conda activate familiarity
 pip install -e .
 ```
 
 ## Usage
 ```python
 import numpy as np
-from familarity import compute_metric
+from familiarity import compute_metric
 train_labels_set = ["person", "location", "building", "eagle", "restaurant", "util"]
 train_probs = [0.4, 0.1, 0.1, 0.1, 0.1, 0.2]
 train_labels = np.random.choice(train_labels_set, size=30000, p=train_probs).tolist()

create_splits_of_verying_difficulty.py

@@ -0,0 +1,180 @@
+import copy
+import json
+from typing import Dict, List
+
+import numpy as np
+import pandas as pd
+import torch
+from datasets import Dataset, DatasetDict
+from sentence_transformers import SentenceTransformer
+from torch.nn.functional import cosine_similarity
+from tqdm import tqdm
+
+
+def create_splits_for_hf_hub(train_dataset: str):
+    # Dataset format should be a list of dictionaries, where each dictionary represents a data point.
+    path_to_train_data = f"path/to/train/{train_dataset}.json"
+    with open(path_to_train_data, "r") as f:
+        data = json.load(f)
+
+    for filter_by in ["entropy", "max"]:
+        dataset_dict = DatasetDict()
+        for setting in ["easy", "medium", "hard"]:
+            new_split = create_splits(
+                data,
+                train_dataset,
+                filter_by=filter_by,
+                setting=setting,
+            )
+
+            hf_format = [convert_to_hf_format(data_point) for data_point in new_split]
+
+            ds = Dataset.from_pandas(pd.DataFrame(data=hf_format))
+            dataset_dict[setting] = ds
+
+        dataset_dict.push_to_hub(f"{train_dataset}_{filter_by}_splits")
+
+
+def convert_to_hf_format(data_point):
+    tags = ["O"] * len(data_point["tokenized_text"])
+    spans = []
+    for ent in data_point["ner"]:
+        start, end, label = ent[0], ent[1], ent[2]
+        spans.append({"start": start, "end": end, "label": label})
+        if start == end:
+            tags[start] = "B-" + label
+        else:
+            try:
+                tags[start] = "B-" + label
+                tags[start + 1 : end + 1] = ["I-" + label] * (end - start)
+            except:
+                pass
+    return {"tokens": data_point["tokenized_text"], "ner_tags": tags, "spans": spans}
+
+
+def create_splits(
+    dataset: List[Dict],
+    dataset_name: str,  # The name of the dataset for which the splits should be created
+    filter_by: str = "entropy",
+    setting: str = "medium",
+):
+    try:
+        df = pd.read_pickle("new_splits.pkl")
+    except:
+        raise FileNotFoundError("Please run the compute_new_splits function first to generate the data.")
+    df = df[(df["train_dataset"] == dataset_name)]
+
+    selected_entity_types = []
+    for benchmark_name in df["eval_dataset"].unique():
+        _df = df[(df["eval_dataset"] == benchmark_name)].copy()
+
+        # The thresholds are dataset specific and may need to be adjusted to account for dataset with different characteristics
+        if filter_by == "entropy":
+            low_threshold = df[filter_by].quantile(0.01)
+            high_threshold = df[filter_by].quantile(0.95)
+        elif filter_by == "max":
+            low_threshold = df[filter_by].quantile(0.05)
+            high_threshold = df[filter_by].quantile(0.99)
+
+        medium_lower_threshold = df[filter_by].quantile(0.495)
+        medium_upper_threshold = df[filter_by].quantile(0.505)
+
+        # Define conditions and choices for categorization
+        conditions = [
+            _df[filter_by] <= low_threshold,  # Bottom
+            _df[filter_by].between(medium_lower_threshold, medium_upper_threshold),  # Middle
+            _df[filter_by] >= high_threshold,  # Top
+        ]
+        choices = ["easy", "medium", "hard"] if filter_by == "entropy" else ["hard", "medium", "easy"]
+
+        # Use np.select to create the new column based on the conditions
+        _df["difficulty"] = np.select(conditions, choices, default="not relevant")
+
+        selected_entity_types.extend(_df[_df["difficulty"] == setting]["entity"].tolist())
+
+    new_dataset = []
+    for dp in tqdm(dataset):
+        matched_entities = [x for x in dp["ner"] if x[-1].lower().strip() in selected_entity_types]
+        if matched_entities:
+            new_np = copy.deepcopy(dp)
+            new_np["ner"] = matched_entities
+            new_dataset.append(new_np)
+
+    return new_dataset
+
+
+def compute_new_splits():
+    # TODO: you need to load the data into two variables: 'benchmarks' and 'training_datasets'.
+    # 'benchmarks' should be a dictionary with the benchmark names as keys and the (list of distinct) entity types as values.
+    # 'training_datasets' should be a dictionary with the training dataset names as keys and the (list of distinct) entity types as values.
+    # We process multiple benchmarks and training datasets in this example, but you can adjust the code to fit your needs.
+    # Further, we stick with the following dataset layout: list of dictionaries, where each dictionary represents a data point.
+    # For example: [{'tokenized_text': [...], 'ner': [(start, end, entity_type), ...]}, ...]
+
+    benchmarks = {}
+    for benchmark_name in ['path/to/eval/dataset1.json', 'path/to/eval/dataset2.json']:
+        # Data loading logic here, e.g.:
+        # tokens, entity_types = load_eval_dataset(benchmark_name)
+        # benchmarks[benchmark_name] = list(entity_types)
+        pass
+
+    training_datasets = {}
+    for train_dataset_name in ['path/to/train/dataset1.json', 'path/to/train/dataset2.json']:
+        # Data loading logic here, e.g.:
+        # tokens, entity_types = load_train_dataset(train_dataset_name)
+        # training_datasets[train_dataset_name] = list(entity_types)
+        pass
+
+    batch_size = 256
+    model = SentenceTransformer("all-mpnet-base-v2").to("cuda")
+    eval_encodings = {}
+    for benchmark_name, entity_types in benchmarks.items():
+        embeddings = model.encode(entity_types, convert_to_tensor=True, device="cuda")
+        eval_encodings[benchmark_name] = embeddings
+
+    results = {}
+    for dataset_name, entity_types in training_datasets.items():
+        for i in tqdm(range(0, len(entity_types), batch_size)):
+            dataset_name = dataset_name.split(".")[0]
+            batch = entity_types[i : i + batch_size]
+            embeddings = model.encode(batch, convert_to_tensor=True, device="cuda")
+            for benchmark_name, eval_embeddings in eval_encodings.items():
+                similarities = torch.clamp(
+                    cosine_similarity(
+                        embeddings.unsqueeze(1),
+                        eval_embeddings.unsqueeze(0),
+                        dim=2,
+                    ),
+                    min=0.0,
+                    max=1.0,
+                )
+                probabilities = torch.nn.functional.softmax(similarities / 0.01, dim=1)
+                entropy_values = -torch.sum(probabilities * torch.log(probabilities + 1e-10), dim=1)
+                max_values, _ = torch.max(similarities, dim=1)
+
+                if dataset_name not in results:
+                    results[dataset_name] = {}
+                if benchmark_name not in results[dataset_name]:
+                    results[dataset_name][benchmark_name] = {}
+
+                for j, entity in enumerate(batch):
+                    if entity not in results[dataset_name][benchmark_name]:
+                        results[dataset_name][benchmark_name][entity] = {}
+                    results[dataset_name][benchmark_name][entity]["entropy"] = entropy_values[j].cpu().numpy().item()
+                    results[dataset_name][benchmark_name][entity]["max"] = max_values[j].cpu().numpy().item()
+
+    entries = []
+    for dataset_name, eval_comparisons in results.items():
+        for benchmark_name, mapping in eval_comparisons.items():
+            for entity, values in mapping.items():
+                entries.append(
+                    {
+                        "entity": entity,
+                        "entropy": values["entropy"],
+                        "max": values["max"],
+                        "eval_dataset": benchmark_name,
+                        "train_dataset": dataset_name,
+                    }
+                )
+    df = pd.DataFrame.from_dict(entries, orient="columns")
+    df.to_pickle("new_splits.pkl")

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
-name = "familarity"
+name = "familiarity"
 dynamic = ["version"]
-description = "Estimating label shift and transfer difficulty using Familarity."
+description = "Estimating label shift and transfer difficulty using Familiarity."
 authors = [{ name = "Jonas Golde", email = "jonas.max.golde@hu-berlin.de" }]
 readme = "README.md"
 requires-python = ">3.8"
@@ -33,7 +33,7 @@ testing = ["pytest"]
 dev = ["black", "isort", "ruff"]
 
 [tool.setuptools]
-packages = ["familarity"]
+packages = ["familiarity"]
 package-dir = { "" = "src" }
 
 [tool.black]

src/familarity/__init__.py

@@ -1,3 +1,3 @@
-from familarity.metric import compute_metric
+from familiarity.metric import compute_metric
 
 __all__ = ["compute_metric"]

src/familarity/embedding_models.py

@@ -6,13 +6,12 @@
 
 import numpy as np
 import torch
+from familiarity.utils import get_device
 from huggingface_hub import repo_exists
 from sentence_transformers import SentenceTransformer
 from tqdm import tqdm
 from transformers import AutoModel, AutoTokenizer
 
-from familarity.utils import get_device
-
 
 class LabelEmbeddingModel(ABC):
     def __init__(self):

src/familarity/metric.py

@@ -5,18 +5,17 @@
 
 import numpy as np
 import pandas as pd
-from tqdm import tqdm
-
-from familarity.embedding_models import LabelEmbeddingModel, load_embedding_model
-from familarity.logger import setup_logger
-from familarity.utils import (
+from familiarity.embedding_models import LabelEmbeddingModel, load_embedding_model
+from familiarity.logger import setup_logger
+from familiarity.utils import (
     clipped_cosine_similarity,
     combine_counters,
     cumsum_until,
     df_to_prettytable,
     iterate_dict_in_batches,
     make_output_path,
 )
+from tqdm import tqdm
 
 
 def compute_embeddings(
@@ -82,29 +81,29 @@ def compute_similarities(
     return similarity_df
 
 
-def compute_familarity(
+def compute_familiarity(
     similarity_df: pd.DataFrame,
     k: int = 1000,
     weighting: str = "zipf",
     output_path: Path = None,
     save_embeddings: bool = False,
 ) -> pd.DataFrame:
-    familarity_data = []
+    familiarity_data = []
 
     for label_test in similarity_df["label_test"].unique():
         test_label_df = similarity_df[similarity_df["label_test"] == label_test]
         test_label_df = test_label_df.sort_values("similarity", ascending=False)
         counts = cumsum_until(test_label_df["count_train"], k)
         sims = test_label_df["similarity"][: len(counts)]
-        familarity = weighted_average(sims, counts, k, weighting=weighting)
-        familarity_data.append({"label": label_test, "familarity": familarity})
+        familiarity = weighted_average(sims, counts, k, weighting=weighting)
+        familiarity_data.append({"label": label_test, "familiarity": familiarity})
 
-    familarity_df = pd.DataFrame(familarity_data)
+    familiarity_df = pd.DataFrame(familiarity_data)
 
     if save_embeddings:
-        familarity_df.to_pickle(output_path / "familarity_df.pkl")
+        familiarity_df.to_pickle(output_path / "familiarity_df.pkl")
 
-    return familarity_df
+    return familiarity_df
 
 
 def weighted_average(
@@ -165,8 +164,8 @@ def compute_metric(
     )
 
     similarity_df = compute_similarities(embedding_df, output_path=output_path, save_embeddings=save_embeddings)
-    familarity_df = compute_familarity(
+    familiarity_df = compute_familiarity(
         similarity_df, k=k, weighting=weighting, output_path=output_path, save_embeddings=save_embeddings
     )
     logger.info("Results:\n")
-    logger.info(df_to_prettytable(familarity_df))
+    logger.info(df_to_prettytable(familiarity_df))

src/familarity/utils.py

@@ -36,7 +36,7 @@ def df_to_prettytable(df: pd.DataFrame) -> PrettyTable:
     for idx, row in df.iterrows():
         table.add_row(row, divider=True if idx + 1 == len(df) else False)
 
-    table.add_row(["Marco-Avg. Familarity", round(df["familarity"].mean().item(), 3)])
+    table.add_row(["Marco-Avg. Familiarity", round(df["familiarity"].mean().item(), 3)])
 
     return table
 

tests/conftest.py

@@ -1,7 +1,6 @@
 import numpy as np
 import pytest
-
-from familarity.embedding_models import LabelEmbeddingModel
+from familiarity.embedding_models import LabelEmbeddingModel
 
 
 @pytest.fixture(scope="module")

tests/test_embedding_models.py

@@ -1,7 +1,6 @@
 import numpy as np
 import pytest
-
-from familarity.embedding_models import (
+from familiarity.embedding_models import (
     FastTextModel,
     GloveModel,
     SentenceTransformerModel,

tests/test_logger.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-from familarity.logger import setup_logger
+from familiarity.logger import setup_logger
 
 
 def test_setup_logger(tmp_path: Path, capsys, caplog):