Support noised text recall and metrics

Add CLI --noise option and support for evaluating recall on noised text embeddings. test_recall now accepts use_noise, can load noised embeddings, and writes results to noise-aware recognition directories/filenames. Improve text noise handling by adding get_random_replacement_char and allowing random-character replacement when character_mask is None. Add Euclidean metric (and mean_euclidean) to text reconstruction metrics and include it in summaries and get_besttext_params outputs. Increase decode_text_embeddings batch size from 32 to 64. Declare ipywidgets (and related jupyter widget packages) in pyproject.toml and update the lockfile to include those packages.

Author: roicort

manage.py

@@ -319,14 +319,24 @@ def test_classifier_command(config):
 # Memory Commands
 @cli.command()
 @click.option("--config", help="YAML configuration.")
-def test_recall(config):
+@click.option(
+    "--noise",
+    is_flag=True,
+    help="Use noised test embeddings if available for text recall evaluation.",
+)
+def test_recall(config, noise):
     "🔍 Test recall with distinct params"
     cfg = load_cli_config(config)
     config_summary(cfg)
 
     if cfg.app.modality == "text":
         from reamember.pipes.text import test_recall
-        test_recall(cfg, device=device, experiments_root=EXPERIMENTS_ROOT)
+        test_recall(
+            cfg,
+            device=device,
+            experiments_root=EXPERIMENTS_ROOT,
+            use_noise=noise,
+        )
         return
     else:
         raise NotImplementedError("Recall testing is only implemented for text modality.")

pyproject.toml

@@ -27,4 +27,5 @@ dependencies = [
     "tqdm>=4.67.3",
     "nltk>=3.9.4",
     "gradio>=6.14.0",
+    "ipywidgets>=8.1.8",
 ]

reamember/pipes/text.py

@@ -3,6 +3,7 @@
 import pandas as pd
 import json
 import math
+import string
 from nltk.metrics import edit_distance
 from tqdm import tqdm
 from plotly import graph_objects as go
@@ -63,7 +64,19 @@ def normalize_noise_level(noise_level):
     return min(noise_level, 1.0)
 
 
-def apply_text_noise(text, noise_level=0.1, character_mask='_'):
+def get_random_replacement_char(original_char):
+    if str(original_char).isupper():
+        alphabet = string.ascii_uppercase
+        original = str(original_char).upper()
+    else:
+        alphabet = string.ascii_lowercase
+        original = str(original_char).lower()
+
+    candidates = [char for char in alphabet if char != original]
+    return str(np.random.choice(candidates))
+
+
+def apply_text_noise(text, noise_level=0.1, character_mask=None):
     text = str(text).strip()
     if not text:
         return text
@@ -95,21 +108,35 @@ def apply_text_noise(text, noise_level=0.1, character_mask='_'):
 
     masked_text = list(text)
     for position in selected_positions:
-        masked_text[position] = character_mask
+        masked_text[position] = (
+            get_random_replacement_char(masked_text[position])
+            if character_mask is None
+            else character_mask
+        )
 
     return ''.join(masked_text)
 
 
-def test_recall(cfg, device, experiments_root):
+def test_recall(cfg, device, experiments_root, use_noise=False):
 
     latent = cfg.neural.latent_dim[0] if isinstance(cfg.neural.latent_dim, ListConfig) else int(cfg.neural.latent_dim)
 
     path = get_experiment_path(cfg, experiments_root, latent)
     embeddings_dataset = load_embeddings_dataset(path, device=device)
+    test_embeddings_dataset = embeddings_dataset
+    recognition_dirname = "recognition"
+
+    if use_noise:
+        test_embeddings_dataset = load_embeddings_dataset(
+            path,
+            device=device,
+            noised=True,
+        )
+        recognition_dirname = f"recognition_noise_{cfg.app.noise}"
 
     # First Experiment: Test Recognition Rates
 
-    recognition_text_path = ensure_directory(path / "recognition")
+    recognition_text_path = ensure_directory(path / recognition_dirname)
 
     quantizer = Quant(
         torch.cat([embeddings_dataset.train.data, embeddings_dataset.test.data], dim=0)
@@ -141,7 +168,10 @@ def test_recall(cfg, device, experiments_root):
         else [float(cfg.memory.sigma)]
     )
 
-    print(f"[INFO] Testing recall with latent={latent}, domains={domains}, sigmas={sigmas}, xis={xis}, iotas={iotas}, kappas={kappas}")
+    print(
+        f"[INFO] Testing recall with latent={latent}, domains={domains}, sigmas={sigmas}, "
+        f"xis={xis}, iotas={iotas}, kappas={kappas}, use_noise={use_noise}"
+    )
 
     train_embeddings = embeddings_dataset.train.data
     split_index = max(1, len(train_embeddings) // 2) # Split train into two halves for seen/unseen evaluation, ensuring at least one sample in the seen half
@@ -152,7 +182,7 @@ def test_recall(cfg, device, experiments_root):
 
     confusion_summaries = [] # Prepare to collect confusion summaries for all domains
     confusion_summary_path = (
-        recognition_text_path / "recognition_confusion_summary.json"
+        recognition_text_path / f"recognition_confusion_summary_{use_noise}.json"
     )
 
     from rich.progress import (
@@ -223,7 +253,7 @@ def test_recall(cfg, device, experiments_root):
                                 eam,
                                 seen_dataset=memory_wrapper.train,
                                 unseen_dataset=memory_wrapper.test,
-                                test_dataset=embeddings_dataset.test,
+                                test_dataset=test_embeddings_dataset.test,
                                 quantizer=quantizer,
                                 batch_size=batch_size,
                             )
@@ -238,15 +268,14 @@ def test_recall(cfg, device, experiments_root):
                                 "xi": xi,
                                 "seen_source": "first_half_of_train",
                                 "unseen_source": "second_half_of_train",
-                                "test_source": "dataset_test_split",
                                 "labels": confusion["labels"],
                                 "matrix": confusion["matrix"].tolist(),
                                 "counts": confusion["counts"],
                                 "rates": confusion["rates"],
                             }
 
                             confusion_path = (
-                                recognition_text_path / f"recognition_confusion_domain_{domain}_sigma_{sigma}_xi_{xi}_iota_{iota}_kappa_{kappa}.json"
+                                recognition_text_path / f"recognition_confusion_domain_{domain}_sigma_{sigma}_xi_{xi}_iota_{iota}_kappa_{kappa}_{test_source}.json"
                             )
 
                             with open(confusion_path, "w", encoding="utf-8") as f_out:
@@ -633,17 +662,18 @@ def text_reconstruction_metrics(
     reconstructed_embeddings = torch.cat(reconstructed_embeddings, dim=0)
     source_embeddings = source_embeddings.detach().cpu()
 
-    source_norm = torch.linalg.norm(source_embeddings, dim=1)
-    reconstructed_norm = torch.linalg.norm(reconstructed_embeddings, dim=1)
-    denom = torch.clamp(source_norm * reconstructed_norm, min=1e-12)
+    source_norm = torch.linalg.norm(source_embeddings, dim=1) # Compute norms for cosine similarity
+    reconstructed_norm = torch.linalg.norm(reconstructed_embeddings, dim=1) # Compute norms for cosine similarity
+    denom = torch.clamp(source_norm * reconstructed_norm, min=1e-12) # Avoid division by zero in cosine similarity calculation
     cosine_scores = (
         torch.sum(source_embeddings * reconstructed_embeddings, dim=1) / denom
-    )
-    l2_scores = torch.linalg.norm(source_embeddings - reconstructed_embeddings, dim=1)
+    ) # Get cosine similarity scores between source and reconstructed embeddings
+    l2_scores = torch.linalg.norm(source_embeddings - reconstructed_embeddings, dim=1) # Get L2 distance scores between source and reconstructed embeddings
+    euclidean_scores = torch.sqrt(torch.sum((source_embeddings - reconstructed_embeddings) ** 2, dim=1)) # Get Euclidean distance scores between source and reconstructed embeddings
 
     samples = []
-    for index, (original, cue, reconstructed, cosine, l2) in enumerate(
-        zip(original_texts, cue_texts, reconstructed_texts, cosine_scores, l2_scores)
+    for index, (original, cue, reconstructed, cosine, l2, euclidean) in enumerate(
+        zip(original_texts, cue_texts, reconstructed_texts, cosine_scores, l2_scores, euclidean_scores)
     ):
         samples.append(
             {
@@ -652,6 +682,7 @@ def text_reconstruction_metrics(
                 "cue": cue,
                 "reconstructed": reconstructed,
                 "cosine": float(cosine.item()),
+                "euclidean": float(euclidean.item()),
                 "l2": float(l2.item()),
                 "edit_distance": int(
                     edit_distance(original.lower(), reconstructed.lower())
@@ -662,6 +693,7 @@ def text_reconstruction_metrics(
     summary = {
         "samples": len(samples),
         "mean_cosine": float(np.mean([item["cosine"] for item in samples])),
+        "mean_euclidean": float(np.mean([item["euclidean"] for item in samples])),
         "mean_l2": float(np.mean([item["l2"] for item in samples])),
         "mean_edit_distance": float(
             np.mean([item["edit_distance"] for item in samples])
@@ -851,6 +883,7 @@ def get_besttext_params(cfg, config, device, EXPERIMENTS_ROOT, use_noise=False):
                                         "recognized": recognized,
                                         "unrecognized": unrecognized,
                                         "mean_cosine": summary["mean_cosine"],
+                                        "mean_euclidean": summary["mean_euclidean"],
                                         "mean_l2": summary["mean_l2"],
                                         "mean_edit_distance": summary["mean_edit_distance"],
                                     }

reamember/utils.py

@@ -177,7 +177,7 @@ def load_model_state(model, model_path, label, device, move_to_device=True):
     return model
 
 
-def decode_text_embeddings(model, embeddings, device=None, batch_size=32):
+def decode_text_embeddings(model, embeddings, device=None, batch_size=64):
     """
     Decode embeddings into text batches using SONAR.
     """