import sys import pandas as pd import numpy as np import math from collections import Counter, defaultdict from typing import List, Any from sklearn.feature_extraction.text import TfidfVectorizer import os import pickle import hashlib import json from gram2vec import vectorizer from openai import OpenAI from openai.lib._pydantic import to_strict_json_schema from pydantic import BaseModel from pydantic import ValidationError import time from utils.llm_feat_utils import generate_feature_spans_cached 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" ZOOM_CACHE = "datasets/zoom_cache/features_cache.json" REGION_CACHE = "datasets/region_cache/regions_cache.pkl" os.makedirs(CACHE_DIR, 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... CACHE_VERSION = 1 # Features to exclude from Gram2Vec outputs EXCLUDED_G2V_FEATURE_PREFIXES = [ 'num_tokens' ] EXCLUDED_G2V_FEATURES = set([ 'num_tokens:num_tokens' ]) class style_analysis_schema(BaseModel): features: list[str] spans: dict[str, dict[str, list[str]]] class FeatureIdentificationSchema(BaseModel): features: list[str] class SpanExtractionSchema(BaseModel): spans: dict[str, dict[str, list[str]]] # {author_name: {feature: [spans]}} def compute_g2v_features(clustered_authors_df: pd.DataFrame, task_authors_df: pd.DataFrame=None, text_clm='fullText') -> pd.DataFrame: """ Computes gram2vec feature vectors for each author and adds them to the DataFrame. This effectively creates a mapping from each author to their vector. """ if task_authors_df is not None: 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"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)}") # Gather the input texts (preserves list-of-strings if any) # If an entry is a list of strings, join; otherwise use the string as-is author_texts = [('\n\n'.join(x) if isinstance(x, list) else x) for x in clustered_authors_df.fullText.tolist()] print(f"Number of author_texts: {len(author_texts)}") # Create a reproducible JSON serialization of the texts serialized = json.dumps({ "col": text_clm, "texts": author_texts }, sort_keys=True, ensure_ascii=False) # Compute MD5 hash digest = hashlib.md5(serialized.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"Cache hit...") with open(cache_path, "rb") as f: clustered_authors_df = pickle.load(f) else: # Else compute and cache g2v_feats_df = vectorizer.from_documents(author_texts, batch_size=8) print(f"Number of g2v features: {len(g2v_feats_df)}") print(f"Number of clustered_authors_df.authorID.tolist(): {len(clustered_authors_df.authorID.tolist())}") print(f"Number of g2v_feats_df.to_numpy().tolist(): {len(g2v_feats_df.to_numpy().tolist())}") ids = clustered_authors_df.authorID.tolist() counter = Counter(ids) duplicates = [k for k, v in counter.items() if v > 1] print(f"Duplicate authorIDs: {duplicates}") print(f"Number of duplicates: {len(ids) - len(set(ids))}") author_to_g2v_feats = {x[0]: x[1] for x in zip(clustered_authors_df.authorID.tolist(), g2v_feats_df.to_numpy().tolist())} print(f"Number of authors with g2v features: {len(author_to_g2v_feats)}") # apply normalization vector_std = np.std(list(author_to_g2v_feats.values()), axis=0) vector_mean = np.mean(list(author_to_g2v_feats.values()), axis=0) vector_std[vector_std == 0] = 1.0 author_to_g2v_feats_z_normalized = {x[0]: (x[1] - vector_mean) / vector_std for x in author_to_g2v_feats.items()} print(f"Number of authors with g2v features normalized: {len(author_to_g2v_feats_z_normalized)}") print(f" len of clustered authors df: {len(clustered_authors_df)}") # Add the vectors as a new column of the DataFrame. clustered_authors_df['g2v_vector'] = [{x[1]: x[0] for x in zip(val, g2v_feats_df.columns.tolist())} for val in author_to_g2v_feats_z_normalized.values()] with open(cache_path, "wb") as f: pickle.dump(clustered_authors_df, f) if task_authors_df is not None: task_authors_df = clustered_authors_df[clustered_authors_df.authorID.isin(task_authors_df.authorID.tolist())] clustered_authors_df = clustered_authors_df[~clustered_authors_df.authorID.isin(task_authors_df.authorID.tolist())] return clustered_authors_df['g2v_vector'].tolist(), task_authors_df['g2v_vector'].tolist() def get_task_authors_from_background_df(background_df): task_authors_df = background_df[background_df.authorID.isin(["Q_author", "a0_author", "a1_author", "a2_author"])] return task_authors_df def instance_to_df(instance, predicted_author=None, ground_truth_author=None): #create a dataframe of the task authors task_authos_df = pd.DataFrame([ {'authorID': 'Mystery author', 'fullText': instance['Q_fullText'], 'predicted': None, 'ground_truth': None}, {'authorID': 'Candidate Author 1', 'fullText': instance['a0_fullText'], 'predicted': int(predicted_author) == 0 if predicted_author is not None else None, 'ground_truth': int(ground_truth_author) == 0 if ground_truth_author is not None else None}, {'authorID': 'Candidate Author 2', 'fullText': instance['a1_fullText'], 'predicted': int(predicted_author) == 1 if predicted_author is not None else None, 'ground_truth': int(ground_truth_author) == 1 if ground_truth_author is not None else None}, {'authorID': 'Candidate Author 3', 'fullText': instance['a2_fullText'], 'predicted': int(predicted_author) == 2 if predicted_author is not None else None, 'ground_truth': int(ground_truth_author) == 2 if ground_truth_author is not None else None} ]) if type(instance['Q_fullText']) == list: task_authos_df = task_authos_df.groupby('authorID').agg({'fullText': lambda x: list(x)}).reset_index() return task_authos_df def generate_style_embedding(background_corpus_df: pd.DataFrame, text_clm: str, model_name: str, dimensionality_reduction: bool = True, dimensions: int = 100) -> pd.DataFrame: """ Generates style embeddings for documents in a background corpus using a specified model. If a row in `text_clm` contains a list of strings, the final embedding for that row is the average of the embeddings of all strings in the list. Args: background_corpus_df (pd.DataFrame): DataFrame containing the corpus. text_clm (str): Name of the column containing the text data (either string or list of strings). model_name (str): Name of the model to use for generating embeddings. Returns: pd.DataFrame: The input DataFrame with a new column for style embeddings. """ from sentence_transformers import SentenceTransformer import torch if model_name not in [ 'gabrielloiseau/LUAR-MUD-sentence-transformers', 'gabrielloiseau/LUAR-CRUD-sentence-transformers', 'miladalsh/light-luar', 'AnnaWegmann/Style-Embedding', ]: print('Model is not supported') return background_corpus_df print(f"Generating style embeddings using {model_name} on column '{text_clm}'...") print(background_corpus_df.fullText.tolist()[:10]) model = SentenceTransformer(model_name) embedding_dim = model.get_sentence_embedding_dimension() # Heuristic to check if the column contains lists of strings by checking the first valid item. # This assumes the column is homogenous. is_list_column = False if not background_corpus_df.empty: # Get the first non-NaN value to inspect its type series_no_na = background_corpus_df[text_clm].dropna() if not series_no_na.empty: first_valid_item = series_no_na.iloc[0] if isinstance(first_valid_item, list): is_list_column = True if is_list_column: # Flatten all texts into a single list for batch processing texts_to_encode = [] row_lengths = [] for text_list in background_corpus_df[text_clm]: # Ensure we handle None, empty lists, or other non-list types gracefully if isinstance(text_list, list) and text_list: texts_to_encode.extend(text_list) row_lengths.append(len(text_list)) else: row_lengths.append(0) if texts_to_encode: all_embeddings = model.encode(texts_to_encode, convert_to_tensor=True, show_progress_bar=True) else: all_embeddings = torch.empty((0, embedding_dim), device=model.device) # Reconstruct and average embeddings for each row final_embeddings = [] current_pos = 0 for length in row_lengths: if length > 0: row_embeddings = all_embeddings[current_pos:current_pos + length] avg_embedding = torch.mean(row_embeddings, dim=0) final_embeddings.append(avg_embedding.cpu().numpy()) current_pos += length else: final_embeddings.append(np.zeros(embedding_dim)) else: # Column contains single strings texts = background_corpus_df[text_clm].fillna("").tolist() # convert_to_tensor=False is faster if we just need numpy arrays embeddings = model.encode(texts, show_progress_bar=True) final_embeddings = list(embeddings) # Apply PCA over the embeddings to reduce the dimentionality if dimensionality_reduction: if len(final_embeddings) > 0 and len(final_embeddings[0]) > dimensions: # Only apply PCA if embeddings exist and dim > dimensions pca = PCA(n_components=dimensions) final_embeddings = pca.fit_transform(final_embeddings) return list(final_embeddings) # ── wrapper with caching ─────────────────────────────────────── def cached_generate_style_embedding(background_corpus_df: pd.DataFrame, text_clm: str, model_name: str, task_authors_df: pd.DataFrame = None) -> pd.DataFrame: """ Wraps `generate_style_embedding`, caching its output in pickle files keyed by an MD5 of (model_name + text list). If the cache exists, loads and returns it instead of recomputing. """ if task_authors_df is not None: 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"Number of background corpus authors: {len(background_corpus_df)}") background_corpus_df = pd.concat([task_authors_df, background_corpus_df]) print(f"Number of authors after concatenation: {len(background_corpus_df)}") # Gather the input texts (preserves list-of-strings if any) texts = background_corpus_df[text_clm].fillna("").tolist() # Create a reproducible JSON serialization of the texts serialized = json.dumps({ "model": model_name, "col": text_clm, "texts": texts }, sort_keys=True, ensure_ascii=False) # Compute MD5 hash digest = hashlib.md5(serialized.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"Cache hit for {model_name} on column '{text_clm}'") print(cache_path) 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}") 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' background_corpus_df[col_name] = task_and_background_embeddings with open(cache_path, "wb") as f: pickle.dump(background_corpus_df, f) if task_authors_df is not None: task_authors_df = background_corpus_df[background_corpus_df.authorID.isin(task_authors_df.authorID.tolist())] background_corpus_df = background_corpus_df[~background_corpus_df.authorID.isin(task_authors_df.authorID.tolist())] 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 # 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(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) else: # Else compute and cache 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) with open(cache_path, "wb") as f: pickle.dump(parsed_response, f) 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""" # Sort author names to ensure consistent key regardless of order sorted_authors = sorted(author_names) key_data = { "authors": sorted_authors, "max_num_feats": max_num_feats } key_string = json.dumps(key_data, sort_keys=True) return hashlib.md5(key_string.encode()).hexdigest() def identify_style_features(author_texts: list[str], author_names: list[str], max_num_feats: int = 5) -> list[str]: cache_key = None if author_names: cache_key = generate_cache_key(author_names, max_num_feats) if os.path.exists(ZOOM_CACHE): with open(ZOOM_CACHE, 'r') as f: cache = json.load(f) else: cache = {} if cache_key in cache: print(f"\nCache hit! Using cached features for authors: {author_names}") return cache[cache_key]["features"] else: print(f"Cache miss. Computing features for authors: {author_names}") client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) prompt = f"""Identify {max_num_feats} writing style features that are commonly between the authors texts. Author Texts: {author_texts} """ def _make_call(): response = client.chat.completions.create( model="gpt-4o-mini", messages=[ {"role": "assistant", "content": "You are a forensic linguist specializing in writing styles."}, {"role": "user", "content": prompt} ], response_format={ "type": "json_schema", "json_schema": { "name": "FeatureIdentificationSchema", "schema": to_strict_json_schema(FeatureIdentificationSchema) } } ) return json.loads(response.choices[0].message.content) features = retry_call(_make_call, FeatureIdentificationSchema).features if cache_key and author_names: cache[cache_key] = { "features": features } # save_cache(cache) with open(ZOOM_CACHE, 'w') as f: json.dump(cache, f, indent=2) print(f"Cached features for authors: {author_names}") def retry_call(call_fn, schema_class, max_attempts=3, wait_sec=2): for attempt in range(max_attempts): try: result = call_fn() # Validate against schema validated = schema_class(**result) return validated except (ValidationError, KeyError, json.JSONDecodeError) as e: print(f"Attempt {attempt + 1} failed with error: {e}") time.sleep(wait_sec) raise RuntimeError("All retry attempts failed for OpenAI call.") def extract_all_spans(authors_df: pd.DataFrame, features: list[str], cluster_label_clm_name: str = 'authorID') -> dict[str, dict[str, list[str]]]: """ For each author, use `generate_feature_spans_cached` to get feature->span mappings. Returns a dict: {author_name: {feature: [spans]}} """ client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) spans_by_author = {} for _, row in authors_df.iterrows(): author_name = str(row[cluster_label_clm_name]) print(author_name) role = f"{author_name}" full_text = row['fullText'] spans = generate_feature_spans_cached(client, full_text, features, role) spans_by_author[author_name] = spans return spans_by_author def compute_clusters_style_representation_3( background_corpus_df: pd.DataFrame, cluster_ids: List[Any], cluster_label_clm_name: str = 'authorID', max_num_feats: int = 20, max_num_documents_per_author=1, max_num_authors=10, max_authors_for_span_extraction=4, min_authors_required: int = 2, top_k: int = 10 ): print(f"Computing style representation for visible clusters: {len(cluster_ids)}") # STEP 1: Identify features on 5 visible authors 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_feat_id = background_corpus_df[background_corpus_df[cluster_label_clm_name].isin(cluster_ids)] author_texts = background_corpus_df_feat_id['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_feat_id[cluster_label_clm_name].tolist()[:max_num_authors] print(f"Number of authors: {len(background_corpus_df_feat_id)}") print(author_names) features = identify_style_features(author_texts, author_names, max_num_feats=max_num_feats) # 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) 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 task_author_names = {'Mystery author', 'Candidate Author 1', 'Candidate Author 2', 'Candidate Author 3'} filtered_task_authors = {author: feat_map for author, feat_map in spans_by_author.items() if author in task_author_names.intersection(set(cluster_ids))} print(filtered_task_authors.keys()) # Build per-author sets of features that have at least one span author_present_feature_sets = [ {feature for feature, spans in feature_map.items() if len(spans) > 0} for _, feature_map in filtered_task_authors.items() ] # If nothing to aggregate (e.g., no task authors in selection), fall back to empty list selected_features_ranked = [] if author_present_feature_sets: coverage_counter = Counter() for present_set in author_present_feature_sets: coverage_counter.update(present_set) # Keep features present in at least `min_authors_required` authors eligible_features = [feat for feat, cnt in coverage_counter.items() if cnt >= int(min_authors_required)] # Preserve original LLM feature ordering as a secondary key where possible feature_original_index = {feat: idx for idx, feat in enumerate(features)} if features else {} selected_features_ranked = sorted( eligible_features, key=lambda f: (-coverage_counter[f], feature_original_index.get(f, 10**9)) )[:int(top_k)] print('filtered set of features (min coverage', min_authors_required, '): ', selected_features_ranked) return { "features": list(selected_features_ranked), "spans": spans_by_author } def compute_clusters_g2v_representation( background_corpus_df: pd.DataFrame, author_ids: List[Any], other_author_ids: List[Any], features_clm_name: str, top_n: int = 10, mode: str = "contrastive", sharedness_method: str = "mean_minus_alpha_std", alpha: float = 0.5 ) -> List[tuple]: # Changed return type to List[tuple] to include scores selected_mask = background_corpus_df['authorID'].isin(author_ids).to_numpy() if not selected_mask.any(): return [] # No documents found for the given cluster_ids selected_feats = background_corpus_df[selected_mask][features_clm_name].tolist() all_g2v_feats = list(selected_feats[0].keys()) # If the user requested a sharedness-based score, compute it and return top-N. if mode == "sharedness": selected_matrix = np.array([list(x.values()) for x in selected_feats], dtype=float) if sharedness_method == "mean": scores = selected_matrix.mean(axis=0) elif sharedness_method in ("mean_minus_alpha_std", "mean-std", "mean_minus_std"): means = selected_matrix.mean(axis=0) stds = selected_matrix.std(axis=0) scores = means - float(alpha) * stds elif sharedness_method == "min": scores = selected_matrix.min(axis=0) else: # Default fallback to mean-minus-alpha*std if unknown method means = selected_matrix.mean(axis=0) stds = selected_matrix.std(axis=0) scores = means - float(alpha) * stds # Rank and return with scores feature_scores = [(feat, score) for feat, score in zip(all_g2v_feats, scores) if score > 0] feature_scores.sort(key=lambda x: x[1], reverse=True) return feature_scores[:top_n] # Return tuples instead of just features # Contrastive mode (default): compute target mean and subtract contrast mean all_g2v_values = np.array([list(x.values()) for x in selected_feats]).mean(axis=0) # If an explicit contrast set is provided, use it; otherwise use everyone outside selection if other_author_ids: explicit_mask = background_corpus_df['authorID'].isin(other_author_ids).to_numpy() # Ensure contrast set is disjoint from the selected set contrast_mask = np.logical_and(explicit_mask, ~selected_mask) else: contrast_mask = ~selected_mask other_selected_feats = background_corpus_df[contrast_mask][features_clm_name].tolist() if len(other_selected_feats) > 0: all_g2v_other_values = np.array([list(x.values()) for x in other_selected_feats]).mean(axis=0) else: # No contrast docs → treat contrast mean as zeros all_g2v_other_values = np.zeros_like(all_g2v_values) final_g2v_feats_values = all_g2v_values - all_g2v_other_values # Compute z-scores for normalization # Get population statistics from all features (both selected and contrast) all_feats = background_corpus_df[features_clm_name].tolist() population_matrix = np.array([list(x.values()) for x in all_feats]) population_mean = population_matrix.mean(axis=0) population_std = population_matrix.std(axis=0) # Avoid division by zero population_std = np.where(population_std == 0, 1, population_std) # Calculate z-scores for the contrastive values z_scores = (final_g2v_feats_values - population_mean) / population_std # Keep only features that have a positive contrastive score top_g2v_feats = sorted( [ (feat, val, z_score) for feat, val, z_score in zip(all_g2v_feats, final_g2v_feats_values, z_scores) if val > 0 and feat not in EXCLUDED_G2V_FEATURES and not any(feat.startswith(p) for p in EXCLUDED_G2V_FEATURE_PREFIXES) ], key=lambda x: -x[1] # Sort by contrastive score ) # Filter in only features that are present in selected_authors selected_authors = {'Mystery author', 'Candidate Author 1', 'Candidate Author 2', 'Candidate Author 3'}.intersection(set(author_ids)) # DEBUG: Print what we're actually working with print(f"[DEBUG] author_ids parameter: {author_ids}") print(f"[DEBUG] Hardcoded selected_authors set: {{'Mystery author', 'Candidate Author 1', 'Candidate Author 2', 'Candidate Author 3'}}") print(f"[DEBUG] Intersection result: {selected_authors}") print(f"[DEBUG] Is selected_authors empty? {len(selected_authors) == 0}") # Filter in only features that are present in selected_authors selected_authors_g2v_data = background_corpus_df[background_corpus_df['authorID'].isin(selected_authors)][features_clm_name].tolist() # print(f"[DEBUG] selected_authors_g2v_data length: {len(selected_authors_g2v_data)}") # print(f"[DEBUG] selected_authors_g2v_data content: {selected_authors_g2v_data}") # Get the actual text documents for the selected authors to verify feature presence selected_authors_docs = background_corpus_df[background_corpus_df['authorID'].isin(selected_authors)]['fullText'].tolist() print(f"[DEBUG] Found {len(selected_authors_docs)} documents for selected authors") # Import find_feature_spans for text-based feature verification try: from gram2vec.feature_locator import find_feature_spans print("[DEBUG] Successfully imported find_feature_spans") except ImportError: print("[WARNING] Could not import find_feature_spans, falling back to vector-based filtering") find_feature_spans = None filtered_features = [] for feature, score, z_score in top_g2v_feats: # DEBUG: Print what we're checking for this feature # print(f"[DEBUG] Checking feature: {feature}") # print(f"[DEBUG] Feature score: {score}, z_score: {z_score}") # Check if the feature has a non-zero value in all of the selected authors feature_presence = [] for i, author_g2v_feats in enumerate(selected_authors_g2v_data): feature_value = author_g2v_feats.get(feature, 0) feature_presence.append(feature_value) # print(f"[DEBUG] Author {i} has feature '{feature}' = {feature_value}") # print(f"[DEBUG] All feature values: {feature_presence}") # print(f"[DEBUG] All values > 0? {[v > 0 for v in feature_presence]}") # print(f"[DEBUG] All values > 0? {all(v > 0 for v in feature_presence)}") # First check: feature must be present in Gram2Vec vectors vector_present = all(author_g2v_feats.get(feature, 0) > 0 for author_g2v_feats in selected_authors_g2v_data) # Second check: feature must be present in actual text documents text_present = True if find_feature_spans and selected_authors_docs: try: # Check if feature appears in at least one document from each selected author for i, doc in enumerate(selected_authors_docs): if isinstance(doc, list): doc_text = '\n\n'.join(doc) else: doc_text = str(doc) spans = find_feature_spans(doc_text, feature) if not spans: # No spans found in this document # print(f"[DEBUG] ✗ Feature '{feature}' not found in document {i} of selected author") text_present = False break # else: # print(f"[DEBUG] ✓ Feature '{feature}' found in document {i} with {len(spans)} spans") except Exception as e: print(f"[WARNING] Error checking text presence for feature '{feature}': {e}") # Fall back to vector-based filtering if text checking fails text_present = vector_present # Feature must pass BOTH checks if vector_present and text_present: filtered_features.append((feature, score, z_score)) # print(f"[DEBUG] ✓ Feature '{feature}' PASSED both vector and text checks") # else: # if not vector_present: # # print(f"[DEBUG] ✗ Feature '{feature}' FAILED vector check") # if not text_present: # # print(f"[DEBUG] ✗ Feature '{feature}' FAILED text check") # # print(f"[DEBUG] ✗ Feature '{feature}' FAILED the filter") print('Filtered G2V features: ', [(f[0], f[2]) for f in filtered_features]) # Print feature names and z-scores return filtered_features[:top_n] # Return tuples with z-scores def compute_task_only_g2v_similarity( background_corpus_df: pd.DataFrame, visible_author_ids: List[Any], features_clm_name: str = 'g2v_vector', top_n: int = 10, require_spans: bool = True ) -> List[tuple]: """ Compute top Gram2Vec features that are shared between the Mystery author and the predicted Candidate author, ignoring background authors and contrast. Selection is limited to task authors within the zoom (i.e., present in `visible_author_ids`). A feature is kept if: - it has a positive value (> 0) for both Mystery and Predicted Candidate, - and (optionally) at least one detected span exists in both authors' texts. Scoring strategy prioritizes features strong in both authors: score = min(mystery_value, predicted_value). Returns a list of (feature_name, score) tuples sorted by score desc, limited to top_n. """ task_names = {'Mystery author', 'Candidate Author 1', 'Candidate Author 2', 'Candidate Author 3'} # Filter to visible task authors is_visible = background_corpus_df['authorID'].isin(visible_author_ids) is_task = background_corpus_df['authorID'].isin(task_names) visible_task_df = background_corpus_df[is_visible & is_task] if visible_task_df.empty: return [] # Identify Mystery author row within the visible set mystery_rows = visible_task_df[visible_task_df['authorID'] == 'Mystery author'] if mystery_rows.empty: # If Mystery is not visible, fall back to using any available Mystery row in the corpus mystery_rows = background_corpus_df[background_corpus_df['authorID'] == 'Mystery author'] if mystery_rows.empty: return [] mystery_row = mystery_rows.iloc[0] # Identify the predicted candidate within the visible set using the 'predicted' flag if present predicted_row = None if 'predicted' in visible_task_df.columns: pred_candidates = visible_task_df[visible_task_df['predicted'] == True] if not pred_candidates.empty: predicted_row = pred_candidates.iloc[0] # If not found in visible, try to find anywhere in the corpus if predicted_row is None and 'predicted' in background_corpus_df.columns: pred_any = background_corpus_df[background_corpus_df['predicted'] == True] # Prefer one that is also a task author pred_any = pred_any[pred_any['authorID'].isin(task_names)] if not pred_any.empty else pred_any if not pred_any.empty: predicted_row = pred_any.iloc[0] # If still not found, we cannot build a pair if predicted_row is None: return [] mystery_vec = mystery_row.get(features_clm_name, {}) predicted_vec = predicted_row.get(features_clm_name, {}) if not isinstance(mystery_vec, dict) or not isinstance(predicted_vec, dict): return [] # Prepare texts for optional span gating def _norm_txt(x): if isinstance(x, list): return '\n\n'.join(x) return str(x) mystery_text = _norm_txt(mystery_row.get('fullText', '')) predicted_text = _norm_txt(predicted_row.get('fullText', '')) try: from gram2vec.feature_locator import find_feature_spans as _find_feature_spans except Exception: _find_feature_spans = None shared_features = [] # Iterate over union of feature keys (both authors share the same feature space in practice) for feature_name in set(list(mystery_vec.keys()) + list(predicted_vec.keys())): # Exclude unwanted features if feature_name in EXCLUDED_G2V_FEATURES or any(feature_name.startswith(p) for p in EXCLUDED_G2V_FEATURE_PREFIXES): continue m_val = float(mystery_vec.get(feature_name, 0.0)) p_val = float(predicted_vec.get(feature_name, 0.0)) # Optional span gate: require at least one span in both texts spans_m = spans_p = None if require_spans and _find_feature_spans is not None: try: spans_m = _find_feature_spans(mystery_text, feature_name) or [] spans_p = _find_feature_spans(predicted_text, feature_name) or [] if len(spans_m) == 0 or len(spans_p) == 0: continue except Exception: # On span errors, skip gating and proceed spans_m = spans_m if spans_m is not None else [] spans_p = spans_p if spans_p is not None else [] # Similarity metric: |m| + |p| - |m - p| score = abs(m_val) + abs(p_val) - abs(m_val - p_val) shared_features.append((feature_name, score, m_val, p_val, len(spans_m) if spans_m is not None else -1, len(spans_p) if spans_p is not None else -1)) # Rank by score desc and return top_n shared_features.sort(key=lambda x: x[1], reverse=True) top = shared_features[:top_n] # Debug print of top-N with values and span counts for presence sanity-check try: print("[DEBUG] Task-only G2V top features (feature, mystery_val, predicted_val, score | spans_mystery, spans_predicted):") for feat_name, sc, m_val, p_val, c_m, c_p in top: print(f" {feat_name} | mystery={m_val:.4f}, predicted={p_val:.4f}, S={sc:.4f} | spans=({c_m}, {c_p})") except Exception: pass return [(f, s) for (f, s, _, _, _, _) in top] 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"]] # 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 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() # 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 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)) 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...") # Extract LUAR embeddings from task authors dataframe mystery_embedding = np.array(task_authors_df.iloc[0][col_name]).reshape(1, -1) candidate_embeddings = np.array([ task_authors_df.iloc[1][col_name], task_authors_df.iloc[2][col_name], task_authors_df.iloc[3][col_name] ]) # Compute cosine similarities similarities = cosine_similarity(mystery_embedding, candidate_embeddings)[0] predicted_author = int(np.argmax(similarities)) print(f"Predicted author is Candidate {predicted_author + 1}") return predicted_author def compute_precomputed_regions(bg_proj, bg_ids, q_proj, c_proj, pred_idx, model_name, n_neighbors=7): """ Compute precomputed regions for mystery author and candidates. Args: bg_proj: (N,2) numpy array with 2D coordinates of background authors bg_ids: list of N author IDs for background authors q_proj: (1,2) numpy array with mystery author coordinates c_proj: (3,2) numpy array with candidate author coordinates n_neighbors: number of closest neighbors to include in each region Returns: dict: mapping region names to bounding boxes and author lists """ print("Computing sugested regions for zoom...") key = f"{hashlib.md5((model_name + str(q_proj.tolist()) + str(c_proj.tolist()) + str(n_neighbors)).encode()).hexdigest()}" if os.path.exists(REGION_CACHE): with open(REGION_CACHE, 'rb') as f: cache = pickle.load(f) else: cache = {} if key in cache: print(f"\nCache hit! Using cached regions.") return cache[key] else: print(f"Cache miss. Computing regions.") regions = {} # All points for distance calculation (mystery + candidates + background) all_points = np.vstack([q_proj.reshape(1, -1), c_proj, bg_proj]) all_ids = ['mystery'] + [f'candidate_{i}' for i in range(3)] + bg_ids def get_region_around_point(center_point, region_name, include_points=None): """Get region around a specific point""" # Ensure center_point is 2D for euclidean_distances if center_point.ndim == 1: center_point = center_point.reshape(1, -1) # Calculate distances from center point to all background authors distances = euclidean_distances(center_point, bg_proj)[0] # Get indices of closest neighbors closest_indices = np.argsort(distances)[:n_neighbors] closest_authors = [bg_ids[i] for i in closest_indices] closest_points = bg_proj[closest_indices] # Include the center point in the region # region_points = np.vstack([center_point.reshape(1, -1), closest_points]) if include_points is not None: region_points = include_points.copy() # Add center point and closest background authors region_points = np.vstack([region_points, center_point, closest_points]) else: # Standard case - just center point and neighbors region_points = np.vstack([center_point, closest_points]) # Calculate bounding box with some padding x_min, x_max = region_points[:, 0].min(), region_points[:, 0].max() y_min, y_max = region_points[:, 1].min(), region_points[:, 1].max() # Add padding (10% of range) x_padding = (x_max - x_min) * 0.1 y_padding = (y_max - y_min) * 0.1 bbox = { 'xaxis': [x_min - x_padding, x_max + x_padding], 'yaxis': [y_min - y_padding, y_max + y_padding] } return { 'bbox': bbox, 'authors': closest_authors, 'center_point': center_point, 'description': f"Region around {region_name} ({len(closest_authors)} closest authors)" } def get_region_between_points(point1, point2, name1, name2): """Get region around the midpoint between two points""" midpoint = (point1 + point2) / 2 region_name = f"{name1} & {name2}" # Include both original points in the region include_points = np.vstack([point1.reshape(1, -1), point2.reshape(1, -1)]) return get_region_around_point(midpoint, region_name, include_points=include_points) # # Region 1: Around mystery author only # regions["Mystery Author Neighborhood"] = get_region_around_point( # q_proj, "Mystery Author" # ) # # Regions 2-4: Around each candidate # for i in range(3): # regions[f"Candidate {i+1} Neighborhood"] = get_region_around_point( # c_proj[i], f"Candidate {i+1}" # ) # Regions 5-7: Between mystery and each candidate for i in range(3): if i == pred_idx: #selecting only mystery and predicted candidate region_name = f"Mystery & Candidate {i+1}" regions[region_name] = get_region_between_points( q_proj, c_proj[i], "Mystery", f"Candidate {i+1}" ) # Regions 8-10: Between candidate pairs candidate_pairs = [(0, 1), (0, 2), (1, 2)] for i, (c1, c2) in enumerate(candidate_pairs): if c1 != pred_idx and c2 != pred_idx: #selecting only the non predicated candidates region_name = f"Candidate {c1+1} & Candidate {c2+1}" regions[region_name] = get_region_between_points( c_proj[c1], c_proj[c2], f"Candidate {c1+1}", f"Candidate {c2+1}" ) # Regions 11-12: Around predicted and ground truth (if different) # This would need predicted_author and ground_truth_author indices # For now, we'll create generic regions # Region 11: Centroid of all task authors (mystery + 3 candidates) # task_centroid = np.mean(np.vstack([q_proj, c_proj]), axis=0) # regions["All Task Authors Centroid"] = get_region_around_point( # task_centroid, "All Task Authors", include_points=np.vstack([q_proj, c_proj]) # ) def serialize_numpy_dtypes(obj): if isinstance(obj, np.ndarray): return obj.tolist() elif isinstance(obj, (np.float32, np.float64)): return float(obj) elif isinstance(obj, (np.int32, np.int64)): return int(obj) elif isinstance(obj, dict): return {key: serialize_numpy_dtypes(value) for key, value in obj.items()} elif isinstance(obj, list): return [serialize_numpy_dtypes(item) for item in obj] else: return obj serializable_regions = serialize_numpy_dtypes(regions) response = json.dumps(serializable_regions, default=str) cache[key] = response with open(REGION_CACHE, 'wb') as f: 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)