updates

app.py

@@ -29,6 +29,7 @@ cfg = load_config()
 download_file_override(cfg.get('background_authors_df_url'), cfg.get('background_authors_df_path'))
 download_file_override(cfg.get('instances_to_explain_url'), cfg.get('instances_to_explain_path'))
 download_file_override(cfg.get('gram2vec_feats_url'), cfg.get('gram2vec_feats_path'))
+download_file_override(cfg.get('gram2vec_cache_url'), cfg.get('gram2vec_cache_path'))
 download_file_override(cfg.get('embeddings_cache_url'), cfg.get('embeddings_cache_path'))
 download_file_override(cfg.get('zoom_cache_url'), cfg.get('zoom_cache_path'))
 download_file_override(cfg.get('region_cache_url'), cfg.get('region_cache_path'))
@@ -142,13 +143,13 @@ def app(share=False):
         # ── Model Selection ─────────────────────────────────
         model_radio = gr.Radio(
             choices=[
+                'AnnaWegmann/Style-Embedding',
                 'gabrielloiseau/LUAR-MUD-sentence-transformers',
                 'gabrielloiseau/LUAR-CRUD-sentence-transformers',
                 'miladalsh/light-luar',
-                'AnnaWegmann/Style-Embedding',
                 'Other'
             ],
-            value='gabrielloiseau/LUAR-MUD-sentence-transformers',
+            value='AnnaWegmann/Style-Embedding',
             label='Choose a Model to inspect'
         )
         print(f"Model choices: {model_radio.choices}")
@@ -168,8 +169,8 @@ def app(share=False):
 
         # ── Task Source Selection ─────────────────────────────────
         task_mode = gr.Radio(
-            choices=["Predefined HRS Task", "Upload Your Own Task"],
-            value="Predefined HRS Task",
+            choices=["Predefined Reddit Task", "Upload Your Own Task"],
+            value="Predefined Reddit Task",
             label="Select Task Source"
         )
 

config/config.yaml

@@ -9,6 +9,9 @@ background_authors_df_url:    "https://huggingface.co/datasets/miladalsh/explana
 gram2vec_feats_path:      "./datasets/gram2vec_feats.csv"
 gram2vec_feats_url:      "https://huggingface.co/datasets/miladalsh/explanation_tool_files/resolve/main/gram2vec_feats.csv?download=true"
 
+gram2vec_cache_url: "https://huggingface.co/datasets/miladalsh/explanation_tool_files/resolve/main/gram2vec_cache.zip?download=true"
+gram2vec_cache_path: "./datasets/gram2vec_cache/"
+
 embeddings_cache_url: "https://huggingface.co/datasets/miladalsh/explanation_tool_files/resolve/main/embeddings_cache.zip?download=true"
 embeddings_cache_path: "./datasets/embeddings_cache/"
 

precompute_caches.py

@@ -8,18 +8,30 @@ import pandas as pd
 from datetime import datetime
 import yaml
 
-# Import your actual modules exactly as app.py does
-from utils.visualizations import get_instances, load_interp_space, trigger_precomputed_region, handle_zoom_with_retries
-from utils.ui import update_task_display
+CONFIG_PATH="config/config.yaml"
 
 def load_config(path="config/config.yaml"):
     with open(path, "r") as f:
         return yaml.safe_load(f)
+    
+# Load configuration and instances EXACTLY like app.py
+cfg = load_config(CONFIG_PATH)
+print(f"Configuration loaded from {CONFIG_PATH}")
+print(f"config : \n{cfg}")
+
+# Import your actual modules exactly as app.py does
+from utils.file_download import download_file_override
+
+download_file_override(cfg.get('background_authors_df_url'), cfg.get('background_authors_df_path'))
+download_file_override(cfg.get('instances_to_explain_url'), cfg.get('instances_to_explain_path'))
+download_file_override(cfg.get('gram2vec_feats_url'), cfg.get('gram2vec_feats_path'))
+
+from utils.visualizations import get_instances, trigger_precomputed_region, handle_zoom_with_retries
+from utils.ui import update_task_display
 
 def precompute_all_caches(
     models_to_test=None,
     instances_to_process=None,
-    config_path="config/config.yaml"
 ):
     """
     Precompute all cache files using the EXACT same methods as app.py.
@@ -34,16 +46,12 @@ def precompute_all_caches(
             'AnnaWegmann/Style-Embedding'
         ]
     
-    print("=" * 60)
+    print("\n\n" + "=" * 60)
     print("CACHE PRECOMPUTATION STARTED")
     print(f"Timestamp: {datetime.now()}")
     print(f"Models to test: {len(models_to_test)}")
     print("=" * 60)
-    
-    # Load configuration and instances EXACTLY like app.py
-    cfg = load_config(config_path)
-    print(f"Configuration loaded from {config_path}")
-    print(f"config : \n{cfg}")
+
     instances, instance_ids = get_instances(cfg['instances_to_explain_path'])
     # interp = load_interp_space(cfg)
     # clustered_authors_df = interp['clustered_authors_df']
@@ -72,7 +80,9 @@ def precompute_all_caches(
         for instance_id in tqdm(instances_to_process, desc=f"Processing instances for {model_name.split('/')[-1]}"):
             current_combination += 1
             try:
-                print(f"\n[{current_combination}/{total_combinations}] Processing Instance {instance_id}")
+                # print(f"\n\n\n[{current_combination}/{total_combinations}] Processing Instance {instance_id}")
+                print(f"\n\n\n\033[1m\033[93m>>> [{current_combination}/{total_combinations}] Processing Instance {instance_id} <<<\033[0m\n")
+
                 
                 # STEP 1: Replicate the exact flow from load_button.click()
                 print("  → Replicating load_button.click() flow...")
@@ -82,7 +92,7 @@ def precompute_all_caches(
                 
                 # Call update_task_display EXACTLY like app.py does
                 task_results = update_task_display(
-                    mode="Predefined HRS Task",  # Always use predefined for caching
+                    mode="Predefined Reddit Task",  # Always use predefined for caching
                     iid=f"Task {instance_id}",
                     instances=instances,
                     background_df=clustered_authors_df,
@@ -137,6 +147,7 @@ def precompute_all_caches(
                 if precomputed_regions_state:
                     regions_dict = ast.literal_eval(precomputed_regions_state)
                     test_regions = list(regions_dict.keys()) 
+                    print(f"    → Found {len(test_regions)} regions to test")
                     
                     for region_name in test_regions:
                         try:
@@ -194,7 +205,7 @@ from utils.visualizations import visualize_clusters_plotly
 
 if __name__ == "__main__":
     # Test with a small subset first
-    instances=[i for i in range(10)]  # First 20 instances for testing
+    instances=[i for i in range(20)]  # First 10 instances for testing
     cache_stats = precompute_all_caches(
         models_to_test=[
              'AnnaWegmann/Style-Embedding'

utils/gram2vec_feat_utils.py

@@ -49,7 +49,17 @@ def get_shorthand(feature_str: str) -> str:
         return None
     if category not in FEATURE_HANDLERS:
         return None
-    code = load_code_map().get(human)
+    code_map = load_code_map()
+    code = code_map.get(human)
+    if code is None:
+        # Try normalizing terminology shown in UI
+        # Convert 'Preposition' phrasing back to 'Adposition' used in the code map
+        human_alt = (human
+                     .replace("Preposition", "Adposition")
+                     .replace("preposition", "adposition")
+                     .replace("Prepositional", "Adpositional")
+                     .replace("prepositional", "adpositional"))
+        code = code_map.get(human_alt)
     if code is None:
         # print(f"Warning: No code found for human-readable feature '{human}'")
         return None  # fallback to the human-readable name
@@ -78,6 +88,14 @@ def get_fullform(shorthand: str) -> str:
     if human is None:
         return None
 
+    # Normalize terminology for UI: prefer "Preposition" over "Adposition"
+    # Also handle potential "adpositional" variants if present
+    human = (human
+             .replace("Adposition", "Preposition")
+             .replace("adposition", "preposition")
+             .replace("Adpositional", "Prepositional")
+             .replace("adpositional", "prepositional"))
+
     return f"{category}:{human}"
 
 def highlight_both_spans(text, llm_spans, gram_spans):
@@ -154,6 +172,8 @@ def show_combined_spans_all(selected_feature_llm, selected_feature_g2v,
         # print(llm_style_feats_analysis)
         print(f"{len(llm_style_feats_analysis['spans'].values())}")
         author_list = list(llm_style_feats_analysis['spans'].values())
+        # print(f"Author list length: {len(author_list)}")
+        # print(f"Author list: {author_list}")
         llm_spans_list = []
         for i, (_, txt) in enumerate(texts):
             print(f"{i}/{len(texts)}")
@@ -169,8 +189,9 @@ def show_combined_spans_all(selected_feature_llm, selected_feature_g2v,
     if selected_feature_g2v and selected_feature_g2v != "None":
         # get gram2vec spans
         gram_spans_list = []
-        # clean the display string and get the feature name without the zscore
-        selected_feature_g2v = selected_feature_g2v.split(" | [Z=")[0].strip()
+        # In case any old label formatting with z-scores leaks through, strip it defensively
+        if "| [Z=" in selected_feature_g2v:
+            selected_feature_g2v = selected_feature_g2v.split(" | [Z=")[0].strip()
         print(f"Selected Gram2Vec feature: {selected_feature_g2v}")
         short = get_shorthand(selected_feature_g2v)
         print(f"short hand: {short}")
@@ -209,12 +230,12 @@ def show_combined_spans_all(selected_feature_llm, selected_feature_g2v,
     bg_start = 4
     bg_indices = list(range(bg_start, len(texts)))
     kept_indices = [i for i in bg_indices if gram_spans_list[i]]
-    print(f"\n---> {kept_indices}")
+    # print(f"\n---> {kept_indices}")
     filtered_texts_bg = [texts[i] for i in kept_indices]
     filtered_llm_bg   = [llm_spans_list[i] for i in kept_indices]
     filtered_gram_bg  = [gram_spans_list[i] for i in kept_indices]
 
-    print(filtered_texts_bg)
+    # print(filtered_texts_bg)
 
     html_background_authors = create_html(
         filtered_texts_bg,
@@ -260,7 +281,7 @@ def get_label(label: str, predicted_author=None, ground_truth_author=None, bg_id
 def create_html(texts, llm_spans_list, gram_spans_list, selected_feature_llm, selected_feature_g2v, short=None, background = False, predicted_author=None, ground_truth_author=None):
     html = []
     for i, (label, txt) in enumerate(texts):
-        print(i, label, txt[:30])
+        # print(i, label, txt[:30])
         label = get_label(label, predicted_author, ground_truth_author,  i) if background else get_label(label, predicted_author, ground_truth_author)
         combined = highlight_both_spans(txt, llm_spans_list[i], gram_spans_list[i])
         notice = ""

utils/interp_space_utils.py

@@ -17,16 +17,20 @@ from pydantic import BaseModel
 from pydantic import ValidationError
 import time
 from utils.llm_feat_utils import generate_feature_spans_cached
+from utils.gram2vec_feat_utils import get_shorthand, get_fullform
+from gram2vec.feature_locator import find_feature_spans
 from collections import Counter
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances
 from sklearn.decomposition import PCA
 
 CACHE_DIR = "datasets/embeddings_cache"
+G2V_CACHE = "datasets/gram2vec_cache"
 ZOOM_CACHE = "datasets/zoom_cache/features_cache.json"
 REGION_CACHE = "datasets/region_cache/regions_cache.pkl"
 SUMMARY_CACHE = "datasets/summary_cache/summaries.json"
 os.makedirs(CACHE_DIR, exist_ok=True)
+os.makedirs(G2V_CACHE, exist_ok=True)
 os.makedirs(os.path.dirname(ZOOM_CACHE), exist_ok=True)
 os.makedirs(os.path.dirname(REGION_CACHE), exist_ok=True)
 # Bump this whenever there is a change etc...
@@ -56,8 +60,8 @@ def compute_g2v_features(clustered_authors_df: pd.DataFrame, task_authors_df: pd
         print (f"concatenating task authors and background corpus authors")
         print(f"Number of task authors: {len(task_authors_df)}")
         print(f"task authors author_ids: {task_authors_df.authorID.tolist()}")
-        print(f"task authors -->")
-        print(task_authors_df)
+        # print(f"task authors -->")
+        # print(task_authors_df)
         print(f"Number of background corpus authors: {len(clustered_authors_df)}")
         clustered_authors_df = pd.concat([task_authors_df, clustered_authors_df])
         print(f"Number of authors after concatenation: {len(clustered_authors_df)}")
@@ -65,10 +69,12 @@ def compute_g2v_features(clustered_authors_df: pd.DataFrame, task_authors_df: pd
     # Gather the input texts (preserves list-of-strings if any)
     #texts = background_corpus_df[text_clm].fillna("").tolist()
     author_texts = ['\n\n'.join(x) for x in clustered_authors_df.fullText.tolist()]
-    print('author_text at 0:{}'.format(author_texts[0]))
+    # print('author_text at 0:{}'.format(author_texts[0]))
     print(f"Number of author_texts: {len(author_texts)}")
 
     # Create a reproducible JSON serialization of the texts
+    # why are g2v features going into a new file inside embeddings_cache?
+    # changed to G2V_CACHE
     serialized = json.dumps({
         "col": text_clm,
         "texts": author_texts
@@ -76,15 +82,20 @@ def compute_g2v_features(clustered_authors_df: pd.DataFrame, task_authors_df: pd
 
     # Compute MD5 hash
     digest = hashlib.md5(serialized.encode("utf-8")).hexdigest()
-    cache_path = os.path.join(CACHE_DIR, f"{digest}.pkl")
+    cache_path = os.path.join(G2V_CACHE, f"{digest}.pkl")
 
     # If cache hit, load and return
     if os.path.exists(cache_path):
-        print(f"Cache hit...")
+        # print(f"Cache hit...")
+        # Making this green to make it stand out from rest of the logs
+        print(f"\n\n\n\033[1m\033[92m>>> Cache hit for {cache_path} <<<\033[0m\n")
         with open(cache_path, "rb") as f:
             clustered_authors_df = pickle.load(f)
     
     else: # Else compute and cache
+        # Making this red to make it stand out from rest of the logs
+        print(f"\n\n\n\033[1m\033[91m>>> Cache miss for {cache_path} => Computing fresh!! <<<\033[0m\n")
+
         g2v_feats_df = vectorizer.from_documents(author_texts, batch_size=8)
 
         print(f"Number of g2v features: {len(g2v_feats_df)}")
@@ -118,6 +129,9 @@ def compute_g2v_features(clustered_authors_df: pd.DataFrame, task_authors_df: pd
         
         with open(cache_path, "wb") as f:
             pickle.dump(clustered_authors_df, f)
+            # Making this green to make it stand out from rest of the logs
+            print(f"\n\n\n\033[1m\033[92m>>> Saved to {cache_path} <<<\033[0m\n")
+            # the file generated here contains g2v + style embeddings. 
     
     if task_authors_df is not None:
         task_authors_df = clustered_authors_df[clustered_authors_df.authorID.isin(task_authors_df.authorID.tolist())]
@@ -268,14 +282,14 @@ def cached_generate_style_embedding(background_corpus_df: pd.DataFrame,
 
     # If cache hit, load and return
     if os.path.exists(cache_path):
-        print(f"Cache hit for {model_name} on column '{text_clm}'")
-        print(cache_path)
+        # Making this green to make it stand out from rest of the logs
+        print(f"\n\n\n\033[1m\033[92m>>> Cache hit for {cache_path} for {model_name} on column '{text_clm} <<<\033[0m\n")
         with open(cache_path, "rb") as f:
             background_corpus_df = pickle.load(f)
 
     else:
         # Otherwise, compute, cache, and return
-        print(f"Computing embeddings for {model_name} on column '{text_clm}', saving to {cache_path}")
+        print(f"\n\n\n\033[1m\033[91m>>> Cache miss for {cache_path} for {model_name} on column '{text_clm} <<<\033[0m\n")
         task_and_background_embeddings = generate_style_embedding(background_corpus_df, text_clm, model_name, dimensionality_reduction=False)
         # Create a clean column name from the model name
         col_name = f'{model_name.split("/")[-1]}_style_embedding'
@@ -283,6 +297,7 @@ def cached_generate_style_embedding(background_corpus_df: pd.DataFrame,
 
         with open(cache_path, "wb") as f:
             pickle.dump(background_corpus_df, f)
+            print(f"\n\n\n\033[1m\033[92m>>> Cache saved for {cache_path} for {model_name} on column '{text_clm} <<<\033[0m\n")
     
     if task_authors_df is not None:
         task_authors_df = background_corpus_df[background_corpus_df.authorID.isin(task_authors_df.authorID.tolist())]
@@ -290,163 +305,167 @@ def cached_generate_style_embedding(background_corpus_df: pd.DataFrame,
 
     return background_corpus_df, task_authors_df
 
-def get_style_feats_distribution(documentIDs, style_feats_dict):
-    style_feats = []
-    for documentId in documentIDs:
-        if documentId not in document_to_style_feats:
-            #print(documentId)
-            continue
-
-        style_feats+= document_to_style_feats[documentId]
-
-    tfidf = [style_feats.count(key) * val for key, val in style_feats_dict.items()]
-
-    return tfidf
-
-def get_cluster_top_feats(style_feats_distribution, style_feats_list, top_k=5):
-    sorted_feats = np.argsort(style_feats_distribution)[::-1]
-    top_feats = [style_feats_list[x] for x in sorted_feats[:top_k] if style_feats_distribution[x] > 0]
-    return top_feats
-
-def compute_clusters_style_representation(
-    background_corpus_df: pd.DataFrame,
-    cluster_ids: List[Any],
-    other_cluster_ids: List[Any],
-    features_clm_name: str,
-    cluster_label_clm_name: str = 'cluster_label',
-    top_n: int = 10
-) -> List[str]:
-    """
-    Given a DataFrame with document IDs, cluster IDs, and feature lists,
-    return the top N features that are most important in the specified `cluster_ids`
-    while having low importance in `other_cluster_ids`.
-    Importance is determined by TF-IDF scores. The final score for a feature is
-    (summed TF-IDF in `cluster_ids`) - (summed TF-IDF in `other_cluster_ids`).
-
-    Parameters:
-    - background_corpus_df: pd.DataFrame. Must contain the columns specified by
-                            `cluster_label_clm_name` and `features_clm_name`.
-                            The column `features_clm_name` should contain lists of strings (features).
-    - cluster_ids: List of cluster IDs for which to find representative features (target clusters).
-    - other_cluster_ids: List of cluster IDs whose features should be down-weighted.
-                         Features prominent in these clusters will have their scores reduced.
-                         Pass an empty list or None if no contrastive clusters are needed.
-    - features_clm_name: The name of the column in `background_corpus_df` that
-                         contains the list of features for each document.
-    - cluster_label_clm_name: The name of the column in `background_corpus_df`
-                              that contains the cluster labels. Defaults to 'cluster_label'.
-    - top_n: Number of top features to return.
-    Returns:
-    - List[str]: A list of feature names. These are up to `top_n` features
-                 ranked by their adjusted TF-IDF scores (score in `cluster_ids`
-                 minus score in `other_cluster_ids`). Only features with a final
-                 adjusted score > 0 are included.
-    """
-
-    assert background_corpus_df[features_clm_name].apply(
-        lambda x: isinstance(x, list) and all(isinstance(feat, str) for feat in x)
-    ).all(), f"Column '{features_clm_name}' must contain lists of strings."
-
-    # Compute TF-IDF on the entire corpus
-    vectorizer = TfidfVectorizer(
-        tokenizer=lambda x: x,
-        preprocessor=lambda x: x,
-        token_pattern=None  # Disable default token pattern, treat items in list as tokens
-    )
-    tfidf_matrix = vectorizer.fit_transform(background_corpus_df[features_clm_name])
-    feature_names = vectorizer.get_feature_names_out()
-
-    # Get boolean mask for documents in selected clusters
-    selected_mask = background_corpus_df[cluster_label_clm_name].isin(cluster_ids).to_numpy()
-
-    if not selected_mask.any():
-        return [] # No documents found for the given cluster_ids
-
-    # Subset the TF-IDF matrix using the boolean mask
-    selected_tfidf = tfidf_matrix[selected_mask]
-
-    # Sum TF-IDF scores across documents for each feature in the target clusters
-    target_feature_scores_sum = selected_tfidf.sum(axis=0).A1  # Convert to 1D array
-
-    # Initialize adjusted scores with target scores
-    adjusted_feature_scores = target_feature_scores_sum.copy()
-
-    # If other_cluster_ids are provided and not empty, subtract their TF-IDF sums
-    if other_cluster_ids: # Checks if the list is not None and not empty
-        other_selected_mask = background_corpus_df[cluster_label_clm_name].isin(other_cluster_ids).to_numpy()
-
-        if other_selected_mask.any():
-            other_selected_tfidf = tfidf_matrix[other_selected_mask]
-            contrast_feature_scores_sum = other_selected_tfidf.sum(axis=0).A1
+# Noticed the following function isnt actually referenced anywhere.
+# def get_style_feats_distribution(documentIDs, style_feats_dict):
+#     style_feats = []
+#     for documentId in documentIDs:
+#         if documentId not in document_to_style_feats:
+#             #print(documentId)
+#             continue
+
+#         style_feats+= document_to_style_feats[documentId]
+
+#     tfidf = [style_feats.count(key) * val for key, val in style_feats_dict.items()]
+
+#     return tfidf
+# 
+# Noticed the following function isnt actually referenced anywhere.
+# def get_cluster_top_feats(style_feats_distribution, style_feats_list, top_k=5):
+#     sorted_feats = np.argsort(style_feats_distribution)[::-1]
+#     top_feats = [style_feats_list[x] for x in sorted_feats[:top_k] if style_feats_distribution[x] > 0]
+#     return top_feats
+
+# Noticed the following function isnt actually referenced anywhere.
+# def compute_clusters_style_representation(
+#     background_corpus_df: pd.DataFrame,
+#     cluster_ids: List[Any],
+#     other_cluster_ids: List[Any],
+#     features_clm_name: str,
+#     cluster_label_clm_name: str = 'cluster_label',
+#     top_n: int = 10
+# ) -> List[str]:
+#     """
+#     Given a DataFrame with document IDs, cluster IDs, and feature lists,
+#     return the top N features that are most important in the specified `cluster_ids`
+#     while having low importance in `other_cluster_ids`.
+#     Importance is determined by TF-IDF scores. The final score for a feature is
+#     (summed TF-IDF in `cluster_ids`) - (summed TF-IDF in `other_cluster_ids`).
+
+#     Parameters:
+#     - background_corpus_df: pd.DataFrame. Must contain the columns specified by
+#                             `cluster_label_clm_name` and `features_clm_name`.
+#                             The column `features_clm_name` should contain lists of strings (features).
+#     - cluster_ids: List of cluster IDs for which to find representative features (target clusters).
+#     - other_cluster_ids: List of cluster IDs whose features should be down-weighted.
+#                          Features prominent in these clusters will have their scores reduced.
+#                          Pass an empty list or None if no contrastive clusters are needed.
+#     - features_clm_name: The name of the column in `background_corpus_df` that
+#                          contains the list of features for each document.
+#     - cluster_label_clm_name: The name of the column in `background_corpus_df`
+#                               that contains the cluster labels. Defaults to 'cluster_label'.
+#     - top_n: Number of top features to return.
+#     Returns:
+#     - List[str]: A list of feature names. These are up to `top_n` features
+#                  ranked by their adjusted TF-IDF scores (score in `cluster_ids`
+#                  minus score in `other_cluster_ids`). Only features with a final
+#                  adjusted score > 0 are included.
+#     """
+
+#     assert background_corpus_df[features_clm_name].apply(
+#         lambda x: isinstance(x, list) and all(isinstance(feat, str) for feat in x)
+#     ).all(), f"Column '{features_clm_name}' must contain lists of strings."
+
+#     # Compute TF-IDF on the entire corpus
+#     vectorizer = TfidfVectorizer(
+#         tokenizer=lambda x: x,
+#         preprocessor=lambda x: x,
+#         token_pattern=None  # Disable default token pattern, treat items in list as tokens
+#     )
+#     tfidf_matrix = vectorizer.fit_transform(background_corpus_df[features_clm_name])
+#     feature_names = vectorizer.get_feature_names_out()
+
+#     # Get boolean mask for documents in selected clusters
+#     selected_mask = background_corpus_df[cluster_label_clm_name].isin(cluster_ids).to_numpy()
+
+#     if not selected_mask.any():
+#         return [] # No documents found for the given cluster_ids
+
+#     # Subset the TF-IDF matrix using the boolean mask
+#     selected_tfidf = tfidf_matrix[selected_mask]
+
+#     # Sum TF-IDF scores across documents for each feature in the target clusters
+#     target_feature_scores_sum = selected_tfidf.sum(axis=0).A1  # Convert to 1D array
+
+#     # Initialize adjusted scores with target scores
+#     adjusted_feature_scores = target_feature_scores_sum.copy()
+
+#     # If other_cluster_ids are provided and not empty, subtract their TF-IDF sums
+#     if other_cluster_ids: # Checks if the list is not None and not empty
+#         other_selected_mask = background_corpus_df[cluster_label_clm_name].isin(other_cluster_ids).to_numpy()
+
+#         if other_selected_mask.any():
+#             other_selected_tfidf = tfidf_matrix[other_selected_mask]
+#             contrast_feature_scores_sum = other_selected_tfidf.sum(axis=0).A1
             
-            # Element-wise subtraction; assumes feature_names aligns for both sums
-            adjusted_feature_scores -= contrast_feature_scores_sum
-
-    # Map scores to feature names
-    feature_score_dict = dict(zip(feature_names, adjusted_feature_scores))
-    # Sort features by score
-    sorted_features = sorted(feature_score_dict.items(), key=lambda item: item[1], reverse=True)
-
-    # Return the names of the top_n features that have a score > 0
-    top_features = [feature for feature, score in sorted_features if score > 0][:top_n]
-
-    return top_features
-
-def compute_clusters_style_representation_2(
-        background_corpus_df: pd.DataFrame,
-        cluster_ids: List[Any],
-        cluster_label_clm_name: str = 'cluster_label',
-        max_num_feats: int = 5,
-        max_num_documents_per_author=3,
-        max_num_authors=5):
-    """
-    Call openAI to analyze the common writing style features of the given list of texts
-    """
-    client = OpenAI(base_url=os.getenv("OPENAI_BASE_URL", None), pi_key=os.getenv("OPENAI_API_KEY"))
-
-    background_corpus_df['fullText'] = background_corpus_df['fullText'].map(lambda x: '\n\n'.join(x[:max_num_documents_per_author]) if isinstance(x, list) else x)
-    background_corpus_df = background_corpus_df[background_corpus_df[cluster_label_clm_name].isin(cluster_ids)]
+#             # Element-wise subtraction; assumes feature_names aligns for both sums
+#             adjusted_feature_scores -= contrast_feature_scores_sum
+
+#     # Map scores to feature names
+#     feature_score_dict = dict(zip(feature_names, adjusted_feature_scores))
+#     # Sort features by score
+#     sorted_features = sorted(feature_score_dict.items(), key=lambda item: item[1], reverse=True)
+
+#     # Return the names of the top_n features that have a score > 0
+#     top_features = [feature for feature, score in sorted_features if score > 0][:top_n]
+
+#     return top_features
+
+# Noticed the following function isnt actually referenced anywhere.
+# def compute_clusters_style_representation_2(
+#         background_corpus_df: pd.DataFrame,
+#         cluster_ids: List[Any],
+#         cluster_label_clm_name: str = 'cluster_label',
+#         max_num_feats: int = 5,
+#         max_num_documents_per_author=3,
+#         max_num_authors=5):
+#     """
+#     Call openAI to analyze the common writing style features of the given list of texts
+#     """
+#     client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+#     background_corpus_df['fullText'] = background_corpus_df['fullText'].map(lambda x: '\n\n'.join(x[:max_num_documents_per_author]) if isinstance(x, list) else x)
+#     background_corpus_df = background_corpus_df[background_corpus_df[cluster_label_clm_name].isin(cluster_ids)]
     
-    author_texts = background_corpus_df['fullText'].tolist()[:max_num_authors]
-    author_texts = "\n\n".join(["""Author {}:\n""".format(i+1) + text for i, text in enumerate(author_texts)])
-    author_names = background_corpus_df[cluster_label_clm_name].tolist()[:max_num_authors]
-    print(f"Number of authors: {len(background_corpus_df)}")
-    print(author_names)
-    print(author_texts)
-    print(f"Number of authors: {len(author_names)}")
-    print(f"Number of authors: {len(author_texts)}")
-    
-    prompt = f"""First identify a list of {max_num_feats} writing style features that are common between the given texts. Second for every author text and style feature, extract all spans that represent the feature. Output for every author all style features with their spans.
-    Author Texts:
-    \"\"\"{author_texts}\"\"\"
-    """
-
-    # Compute MD5 hash
-    digest = hashlib.md5(prompt.encode("utf-8")).hexdigest()
-    cache_path = os.path.join(CACHE_DIR, f"{digest}.pkl")
-
-    # If cache hit, load and return
-    if os.path.exists(cache_path):
-        print(f"Loading authors writing style from cache ...")
-        with open(cache_path, "rb") as f:
-            parsed_response = pickle.load(f)
+#     author_texts = background_corpus_df['fullText'].tolist()[:max_num_authors]
+#     author_texts = "\n\n".join(["""Author {}:\n""".format(i+1) + text for i, text in enumerate(author_texts)])
+#     author_names = background_corpus_df[cluster_label_clm_name].tolist()[:max_num_authors]
+#     print(f"Number of authors: {len(background_corpus_df)}")
+#     print(author_names)
+#     print(author_texts)
+#     print(f"Number of authors: {len(author_names)}")
+#     print(f"Number of authors: {len(author_texts)}")
     
-    else: # Else compute and cache
+#     prompt = f"""First identify a list of {max_num_feats} writing style features that are common between the given texts. Second for every author text and style feature, extract all spans that represent the feature. Output for every author all style features with their spans.
+#     Author Texts:
+#     \"\"\"{author_texts}\"\"\"
+#     """
+
+#     # Compute MD5 hash
+#     digest = hashlib.md5(prompt.encode("utf-8")).hexdigest()
+#     cache_path = os.path.join(CACHE_DIR, f"{digest}.pkl")
+
+#     # If cache hit, load and return
+#     if os.path.exists(cache_path):
+#         print(f"Loading authors writing style from cache ...")
+#         with open(cache_path, "rb") as f:
+#             parsed_response = pickle.load(f)
+    
+#     else: # Else compute and cache
 
-        response = client.chat.completions.create(
-            model="gpt-4o",
-            messages=[
-                {"role":"assistant","content":"You are a forensic linguistic who knows how to analyze similarites in writing styles."},
-                {"role":"user","content":prompt}],
-            response_format={"type": "json_schema", "json_schema": {"name": "style_analysis_schema", "schema":  to_strict_json_schema(style_analysis_schema)}}
-        )
+#         response = client.chat.completions.create(
+#             model="gpt-4o-mini",
+#             messages=[
+#                 {"role":"assistant","content":"You are a forensic linguistic who knows how to analyze similarites in writing styles."},
+#                 {"role":"user","content":prompt}],
+#             response_format={"type": "json_schema", "json_schema": {"name": "style_analysis_schema", "schema":  to_strict_json_schema(style_analysis_schema)}}
+#         )
 
-        parsed_response = json.loads(response.choices[0].message.content)
+#         parsed_response = json.loads(response.choices[0].message.content)
 
-        with open(cache_path, "wb") as f:
-            pickle.dump(parsed_response, f)
+#         with open(cache_path, "wb") as f:
+#             pickle.dump(parsed_response, f)
 
-    return parsed_response 
+#     return parsed_response 
 
 def generate_cache_key(author_names: List[str], max_num_feats: int) -> str:
     """Generate a unique cache key based on author names and max features"""
@@ -472,10 +491,11 @@ def identify_style_features(author_texts: list[str], author_names: list[str], ma
 
         if cache_key in cache:
             print(f"\nCache hit! Using cached features for authors: {author_names}")
+            print(f"\n\n\n\033[1m\033[92m>>> Cache hit for {cache_key} in {ZOOM_CACHE} <<<\033[0m\n")
             return cache[cache_key]["features"]
         else:
-            print(f"Cache miss. Computing features for authors: {author_names}")
-    
+            print(f"\n\n\n\033[1m\033[91m>>> Cache miss for {cache_key} in {ZOOM_CACHE} \nComputing features for authors: {author_names}<<<\033[0m\n")
+
     client = OpenAI(base_url=os.getenv("OPENAI_BASE_URL", None), api_key=os.getenv("OPENAI_API_KEY"))
     prompt = f"""Identify {max_num_feats} writing style features that are common between the authors texts.
     Author Texts:
@@ -483,9 +503,9 @@ def identify_style_features(author_texts: list[str], author_names: list[str], ma
     {author_texts}
     """
 
-    print('==================>>>>>>>>>>')
-    print(prompt)
-    print('==================>>>>>>>>>>')
+    # print('==================>>>>>>>>>>')
+    # print(prompt)
+    # print('==================>>>>>>>>>>')
     def _make_call():
         response = client.chat.completions.create(
             model="gpt-4o",
@@ -512,6 +532,8 @@ def identify_style_features(author_texts: list[str], author_names: list[str], ma
         # save_cache(cache)
         with open(ZOOM_CACHE, 'w') as f:
             json.dump(cache, f, indent=2)
+            print(f"\n\n\n\033[1m\033[92m>>> Cache saved for {cache_key} in {ZOOM_CACHE}<<<\033[0m\n")
+            
 
         print(f"Cached features for authors: {author_names}")
     
@@ -540,7 +562,7 @@ def extract_all_spans(authors_df: pd.DataFrame, features: list[str], cluster_lab
 
     for _, row in authors_df.iterrows():
         author_name = str(row[cluster_label_clm_name])
-        print(author_name)
+        # print(author_name)
         role = f"{author_name}"
         full_text = row['fullText']
         spans = generate_feature_spans_cached(client, full_text, features, role)
@@ -569,18 +591,18 @@ def compute_clusters_style_representation_3(
     author_texts = "\n\n".join(["""Author {}:\n""".format(i+1) + text for i, text in enumerate(author_texts)])
     author_names = background_corpus_df_feat_id[cluster_label_clm_name].tolist()[:max_num_authors]
     print(f"Number of authors: {len(background_corpus_df_feat_id)}")
-    print(author_names)
+    # print(author_names)
     features = identify_style_features(author_texts, author_names, max_num_feats=max_num_feats)
 
     if return_only_feats:
         return features
     
-    print("Features: ", features)
+    #print("Features: ", features)
     # STEP 2: Prepare author pool for span extraction
     span_df = background_corpus_df.iloc[:max_authors_for_span_extraction]
     author_names = span_df[cluster_label_clm_name].tolist()[:max_authors_for_span_extraction]
     print(f"Number of authors for span detection : {len(span_df)}")
-    print(author_names)
+    # print(author_names)
     spans_by_author = extract_all_spans(span_df, features, cluster_label_clm_name)
     
     # Filter-in only task authors that are part of the current selection
@@ -597,7 +619,7 @@ def compute_clusters_style_representation_3(
             for feature, spans in feature_map.items():
                 if spans:
                     feature_importance[feature] -= len(spans)
-    print(feature_importance)
+    # print(feature_importance)
     selected_features_ranked = sorted(feature_importance, key=lambda f: -feature_importance[f])[:int(top_k)]
 
     #print('filtered set of features (min coverage', len(author_present_feature_sets), '): ', selected_features_ranked)
@@ -716,6 +738,7 @@ def compute_clusters_g2v_representation(
     other_author_ids: List[Any],
     features_clm_name: str,
     top_n: int = 10,
+    max_candidates_for_span_sorting: int = 50,
 ) -> List[tuple]:  # Changed return type to List[tuple] to include scores
 
     # 1) Identify selected authors in the zoom region
@@ -749,67 +772,114 @@ def compute_clusters_g2v_representation(
     # 5) Rank features by mean z-score, keep positives only
     feature_scores = [(feat, float(score)) for feat, score in zip(all_features, selected_mean) if score > 0]
     feature_scores.sort(key=lambda x: x[1], reverse=True)
+    
+    # 6) Extract top candidates for span-based sorting
+    candidate_features = feature_scores[:max_candidates_for_span_sorting]
+    
+    # 7) Extract spans for task authors to sort by frequency
+    task_author_names = {'Mystery author', 'Candidate Author 1', 'Candidate Author 2', 'Candidate Author 3'}
+    task_authors_in_selection = [aid for aid in author_ids if aid in task_author_names]
+    
+    if not task_authors_in_selection:
+        # If no task authors in selection, just return the z-score sorted features
+        print("[INFO] No task authors in selection, returning z-score sorted features")
+        return feature_scores[:top_n]
+    
+    # Get task author data
+    task_authors_df = background_corpus_df[background_corpus_df['authorID'].isin(task_authors_in_selection)]
+    
+    # Count spans for each feature across task authors
+    feature_span_counts = {}
+    for feat_shorthand, z_score in candidate_features:
+        span_count = 0
+        
+        # Convert shorthand to human-readable for display (if needed)
+        # Note: features in gram2vec dict are in shorthand format like "pos_unigrams:ADJ"
+        
+        for _, author_row in task_authors_df.iterrows():
+            author_text = author_row['fullText']
+            if isinstance(author_text, list):
+                author_text = '\n\n'.join(author_text)
+            
+            try:
+                # find_feature_spans expects shorthand format like "pos_unigrams:ADJ"
+                spans = find_feature_spans(author_text, feat_shorthand)
+                span_count += len(spans)
+            except Exception as e:
+                # If span extraction fails, continue with 0 spans for this author
+                pass
+        
+        feature_span_counts[feat_shorthand] = span_count
+    
+    # 8) Sort features by span frequency, then by z-score as tiebreaker
+    sorted_by_spans = sorted(
+        candidate_features,
+        key=lambda x: (-feature_span_counts.get(x[0], 0), -x[1])
+    )
+    
+    # print(f"[INFO] Sorted gram2vec features by span frequency: {[(f, feature_span_counts.get(f, 0), z) for f, z in sorted_by_spans[:top_n]]}")
+    
+    return sorted_by_spans[:top_n]
 
-    return feature_scores[:top_n]
-
-def generate_interpretable_space_representation(interp_space_path, styles_df_path, feat_clm, output_clm, num_feats=5):
+# Noticed the following function isnt actually referenced anywhere.
+# def generate_interpretable_space_representation(interp_space_path, styles_df_path, feat_clm, output_clm, num_feats=5):
     
-    styles_df = pd.read_csv(styles_df_path)[[feat_clm, "documentID"]]
+#     styles_df = pd.read_csv(styles_df_path)[[feat_clm, "documentID"]]
 
-    # A dictionary of style features and their IDF
-    style_feats_agg_df = styles_df.groupby(feat_clm).agg({'documentID': lambda x : len(list(x))}).reset_index()
-    style_feats_agg_df['document_freq'] = style_feats_agg_df.documentID
-    style_to_feats_dfreq = {x[0]: math.log(styles_df.documentID.nunique()/x[1]) for x in zip(style_feats_agg_df[feat_clm].tolist(), style_feats_agg_df.document_freq.tolist())}
+#     # A dictionary of style features and their IDF
+#     style_feats_agg_df = styles_df.groupby(feat_clm).agg({'documentID': lambda x : len(list(x))}).reset_index()
+#     style_feats_agg_df['document_freq'] = style_feats_agg_df.documentID
+#     style_to_feats_dfreq = {x[0]: math.log(styles_df.documentID.nunique()/x[1]) for x in zip(style_feats_agg_df[feat_clm].tolist(), style_feats_agg_df.document_freq.tolist())}
     
-    # A list of style features we work with
-    style_feats_list = style_feats_agg_df[feat_clm].tolist()
-    print('Number of style feats ', len(style_feats_list))
+#     # A list of style features we work with
+#     style_feats_list = style_feats_agg_df[feat_clm].tolist()
+#     print('Number of style feats ', len(style_feats_list))
     
-    # A list of documents and what list of style features each has
-    doc_style_agg_df     = styles_df.groupby('documentID').agg({feat_clm: lambda x : list(x)}).reset_index()
-    document_to_feats_dict = {x[0]: x[1] for x in zip(doc_style_agg_df.documentID.tolist(), doc_style_agg_df[feat_clm].tolist())}
+#     # A list of documents and what list of style features each has
+#     doc_style_agg_df     = styles_df.groupby('documentID').agg({feat_clm: lambda x : list(x)}).reset_index()
+#     document_to_feats_dict = {x[0]: x[1] for x in zip(doc_style_agg_df.documentID.tolist(), doc_style_agg_df[feat_clm].tolist())}
     
     
 
-    # Load the clustering information
-    df = pd.read_pickle(interp_space_path)
-    df = df[df.cluster_label != -1]
-    # A cluster to list of documents
-    clusterd_df = df.groupby('cluster_label').agg({
-        'documentID': lambda x: [d_id for doc_ids in x for d_id in doc_ids]
-    }).reset_index()
+#     # Load the clustering information
+#     df = pd.read_pickle(interp_space_path)
+#     df = df[df.cluster_label != -1]
+#     # A cluster to list of documents
+#     clusterd_df = df.groupby('cluster_label').agg({
+#         'documentID': lambda x: [d_id for doc_ids in x for d_id in doc_ids]
+#     }).reset_index()
     
-    # Filter-in only documents that has a style description
-    clusterd_df['documentID'] = clusterd_df.documentID.apply(lambda documentIDs: [documentID for documentID in documentIDs if documentID in document_to_feats_dict])
-    # Map from cluster label to list of features through the document information
-    clusterd_df[feat_clm] = clusterd_df.documentID.apply(lambda doc_ids: [f for d_id in doc_ids for f in document_to_feats_dict[d_id]])
-
-    def compute_tfidf(row):
-        style_counts = Counter(row[feat_clm])
-        total_num_styles = sum(style_counts.values())
-        #print(style_counts, total_num_styles)
-        style_distribution = {
-            style: math.log(1+count) * style_to_feats_dfreq[style] if style in style_to_feats_dfreq else 0 for style, count in style_counts.items()
-        } #TF-IDF
+#     # Filter-in only documents that has a style description
+#     clusterd_df['documentID'] = clusterd_df.documentID.apply(lambda documentIDs: [documentID for documentID in documentIDs if documentID in document_to_feats_dict])
+#     # Map from cluster label to list of features through the document information
+#     clusterd_df[feat_clm] = clusterd_df.documentID.apply(lambda doc_ids: [f for d_id in doc_ids for f in document_to_feats_dict[d_id]])
+
+#     def compute_tfidf(row):
+#         style_counts = Counter(row[feat_clm])
+#         total_num_styles = sum(style_counts.values())
+#         #print(style_counts, total_num_styles)
+#         style_distribution = {
+#             style: math.log(1+count) * style_to_feats_dfreq[style] if style in style_to_feats_dfreq else 0 for style, count in style_counts.items()
+#         } #TF-IDF
         
-        return style_distribution
+#         return style_distribution
 
-    def create_tfidf_rep(tfidf_dist, num_feats):
-        style_feats = sorted(tfidf_dist.items(), key=lambda x: -x[1])
-        top_k_feats = [x[0] for x in style_feats[:num_feats] if str(x[0]) != 'nan']
-        return top_k_feats
+#     def create_tfidf_rep(tfidf_dist, num_feats):
+#         style_feats = sorted(tfidf_dist.items(), key=lambda x: -x[1])
+#         top_k_feats = [x[0] for x in style_feats[:num_feats] if str(x[0]) != 'nan']
+#         return top_k_feats
 
-    clusterd_df[output_clm +'_dist'] = clusterd_df.apply(lambda row: compute_tfidf(row), axis=1)
-    clusterd_df[output_clm]         = clusterd_df[output_clm +'_dist'].apply(lambda dist: create_tfidf_rep(dist, num_feats))
+#     clusterd_df[output_clm +'_dist'] = clusterd_df.apply(lambda row: compute_tfidf(row), axis=1)
+#     clusterd_df[output_clm]         = clusterd_df[output_clm +'_dist'].apply(lambda dist: create_tfidf_rep(dist, num_feats))
 
         
-    return clusterd_df
+#     return clusterd_df
 
 def compute_predicted_author(task_authors_df: pd.DataFrame, col_name: str) -> int:
     """
     Computes the predicted author based on the style features.
     """
-    print("Computing predicted author using LUAR-MUD-style embeddings...")
+    print("Computing predicted author using embeddings...")
 
     # Extract LUAR embeddings from task authors dataframe
     mystery_embedding = np.array(task_authors_df.iloc[0][col_name]).reshape(1, -1)
@@ -850,11 +920,11 @@ def compute_precomputed_regions(bg_proj, bg_ids, q_proj, c_proj, pred_idx, model
     else:
         cache = {}
     if key in cache:
-        print(f"\nCache hit! Using cached regions.")
+        print(f"\n\n\n\033[1m\033[92m>>> Cache hit for {key} in {REGION_CACHE}: Using cached regions<<<\033[0m\n")
         return cache[key]
     else:
-        print(f"Cache miss. Computing regions.")
-    
+        print(f"\n\n\n\033[1m\033[91m>>> Cache miss for {key} in {REGION_CACHE}: Computing Regions<<<\033[0m\n")
+
     regions = {}
     
     # All points for distance calculation (mystery + candidates + background)
@@ -969,22 +1039,23 @@ def compute_precomputed_regions(bg_proj, bg_ids, q_proj, c_proj, pred_idx, model
     response = json.dumps(serializable_regions, default=str)
     cache[key] = response
     with open(REGION_CACHE, 'wb') as f: 
+        print(f"\n\n\n\033[1m\033[92m>>> Cache saved for {key} in {REGION_CACHE} <<<\033[0m\n")
         pickle.dump(cache, f)
 
     return response
 
-if __name__ == "__main__":
-    background_corpus = pd.read_pickle('../datasets/luar_interp_space_cluster_19/train_authors.pkl')
-    print(background_corpus.columns)
-    print(background_corpus[['authorID', 'fullText', 'cluster_label']].head())
-    # # Example: Find features for clusters [2,3,4] that are NOT prominent in cluster [1]
-    # feats = compute_clusters_style_representation(
-    #     background_corpus_df=background_corpus,
-    #     cluster_ids=['00005a5c-5c06-3a36-37f9-53c6422a31d8',],
-    #     other_cluster_ids=[], # Pass the contrastive cluster IDs here
-    #     cluster_label_clm_name='authorID',
-    #     features_clm_name='final_attribute_name'
-    # )
-    # print(feats)
-    generate_style_embedding(background_corpus, 'fullText', 'AnnaWegmann/Style-Embedding')
-    print(background_corpus.columns)
+# if __name__ == "__main__":
+#     background_corpus = pd.read_pickle('../datasets/luar_interp_space_cluster_19/train_authors.pkl')
+#     print(background_corpus.columns)
+#     print(background_corpus[['authorID', 'fullText', 'cluster_label']].head())
+#     # # Example: Find features for clusters [2,3,4] that are NOT prominent in cluster [1]
+#     # feats = compute_clusters_style_representation(
+#     #     background_corpus_df=background_corpus,
+#     #     cluster_ids=['00005a5c-5c06-3a36-37f9-53c6422a31d8',],
+#     #     other_cluster_ids=[], # Pass the contrastive cluster IDs here
+#     #     cluster_label_clm_name='authorID',
+#     #     features_clm_name='final_attribute_name'
+#     # )
+#     # print(feats)
+#     generate_style_embedding(background_corpus, 'fullText', 'AnnaWegmann/Style-Embedding')
+#     print(background_corpus.columns)

utils/llm_feat_utils.py

@@ -32,19 +32,20 @@ def generate_feature_spans(client, text: str, features: list[str]) -> str:
     """
     Call to OpenAI to extract spans. Returns a JSON string.
     """
+    # For some of the longer features, openai client was truncating the feature names, resulting in downstream errors.
+    # Adding structured JSON template to ensure all features are included properly.
+    features_json_template = {feature: [] for feature in features}
     prompt = f"""You are a linguistic specialist. Given a writing sample and a list of descriptive features, identify the exact text spans that demonstrate each feature.
     
     Important:
     - The headers like "Document 1:" etc are NOT part of the original text — ignore them.
     - For each feature, even if there is no match, return an empty list.
     - Only return exact phrases from the text.
+    - Use the EXACT feature names as JSON keys - do not paraphrase or shorten them.
 
-    Respond in JSON format like:
-    {{
-      "feature1": ["span1", "span2"],
-      "feature2": [],
-      …
-    }}
+
+    Respond in this EXACT JSON format (use these exact keys, populate the lists with the extracted text spans):
+    {json.dumps(features_json_template, indent=2)}
 
     Text:
     \"\"\"{text}\"\"\"
@@ -52,9 +53,9 @@ def generate_feature_spans(client, text: str, features: list[str]) -> str:
     Style Features:
     {features}
     """
-    print('==================>>>>>>>>>>')
-    print(prompt)
-    print('==================>>>>>>>>>>')
+    # print('==================>>>>>>>>>>')
+    # print(prompt)
+    # print('==================>>>>>>>>>>')
     response = client.chat.completions.create(
         model="gpt-4o",
         messages=[{"role":"user","content":prompt}]
@@ -71,8 +72,14 @@ def generate_feature_spans_with_retries(client, text: str, features: list[str])
     for attempt in range(MAX_ATTEMPTS):
         try:
             response_str = generate_feature_spans(client, text, features)
-            print(response_str)
+            # print(response_str)
             result = json.loads(response_str)
+            # Additional check to ensure all requested features are present in the response correctly
+            if result.keys() != set(features):
+                print("Response keys do not match requested features. Retrying!")
+                response_str = generate_feature_spans(client, text, features)
+                # print(response_str)
+                result = json.loads(response_str)
             return result
         except (JSONDecodeError, ValueError) as e:
             print(f"Attempt {attempt+1} failed: {e}")
@@ -116,7 +123,13 @@ def generate_feature_spans_cached(client, text: str, features: list[str], role:
         if h in cache:
             # print(f"Found feature: {feat}")
             found_feats_count += 1
-            result[feat] = cache[h]["spans"]
+            if cache[h]["spans"] is None:
+                print(f"Missing feature: {feat}")
+                missing_feats_count += 1
+                missing_feats.append(feat)
+            else:
+                result[feat] = cache[h]["spans"]
+
         else:
             # print(f"Missing feature: {feat}")
             missing_feats_count += 1

utils/ui.py

@@ -81,14 +81,14 @@ def read_txt(f):
 def toggle_task(mode):
     print(mode)
     return (
-        gr.update(visible=(mode == "Predefined HRS Task")),
+        gr.update(visible=(mode == "Predefined Reddit Task")),
         gr.update(visible=(mode == "Upload Your Own Task"))
     )
 
 # Update displayed texts based on mode
 def update_task_display(mode, iid, instances, background_df, mystery_file, cand1_file, cand2_file, cand3_file, true_author, model_radio, custom_model_input):
     model_name = model_radio if model_radio != "Other" else custom_model_input
-    if mode == "Predefined HRS Task":
+    if mode == "Predefined Reddit Task":
         iid = int(iid.replace('Task ', ''))
         data = instances[iid]
         ground_truth_author = 100#data['gt_idx']

utils/visualizations.py

@@ -309,7 +309,7 @@ def handle_zoom(event_json, bg_proj, bg_lbls, clustered_authors_df, task_authors
 
     task_texts = [_to_text(x) for x in task_only_df['fullText'].tolist()]
 
-    print(f"task_texts: {task_texts}")
+    print(f"len task_texts: {len(task_texts)}")
     filtered_g2v_feats = []
     for feat in g2v_feats:
         try:
@@ -333,7 +333,7 @@ def handle_zoom(event_json, bg_proj, bg_lbls, clustered_authors_df, task_authors
     HR_g2v_list = []
     for feat in filtered_g2v_feats:
         HR_g2v = get_fullform(feat[0])
-        print(f"\n\n feat: {feat} ---> Human Readable: {HR_g2v}")
+        # print(f"\n\n feat: {feat} ---> Human Readable: {HR_g2v}")
         if HR_g2v is None:
             print(f"Skipping Gram2Vec feature without human readable form: {feat}")
         else:
@@ -342,11 +342,11 @@ def handle_zoom(event_json, bg_proj, bg_lbls, clustered_authors_df, task_authors
     HR_g2v_list = [("None", None)] + HR_g2v_list
 
     print(f"[INFO] Found {len(llm_feats)} LLM features and {len(g2v_feats)} Gram2Vec features in the zoomed region.")   
-    print(f"[INFO] unfiltered g2v features: {g2v_feats}")
+    # print(f"[INFO] unfiltered g2v features: {g2v_feats}")
 
     print(f"[INFO] LLM features: {llm_feats}")
     HR_g2v_list, _ = format_g2v_features_for_display(HR_g2v_list)
-    print(f"[INFO] Gram2Vec features: {HR_g2v_list}")
+    # print(f"[INFO] Gram2Vec features: {HR_g2v_list}")
 
     return (
         gr.update(choices=llm_feats, value=llm_feats[0]),
@@ -386,7 +386,7 @@ def handle_zoom_with_retries(event_json, bg_proj, bg_lbls, clustered_authors_df,
 def visualize_clusters_plotly(iid, cfg, instances, model_radio, custom_model_input, task_authors_df, background_authors_embeddings_df, pred_idx=None, gt_idx=None):
     model_name = model_radio if model_radio != "Other" else custom_model_input
     embedding_col_name = f'{model_name.split("/")[-1]}_style_embedding'
-    print(background_authors_embeddings_df.columns)
+    # print(background_authors_embeddings_df.columns)
     print("Generating cluster visualization")
     iid = int(iid)
     #interp      = load_interp_space(cfg)