Merge pull request #18 from jeremymanning/main

Implement analysis variants feature (Issues #13, #16, #17)

Author: jeremymanning

README.md

@@ -143,6 +143,127 @@ python generate_figures.py --list
 
 **Note**: The t-test calculations (Figure 2) take approximately 2-3 minutes due to statistical computations across all epochs and authors.
 
+## Analysis Variants
+
+The project supports three linguistic analysis variants to understand what stylistic features models learn:
+
+### Content-Only Variant
+Masks function words with `<FUNC>` token, preserving only content words (nouns, verbs, adjectives, etc.)
+- **Tests:** Whether models distinguish authors based on vocabulary and word choice
+- **Example transformation:**
+  - Original: "The quick brown fox jumps over the lazy dog"
+  - Transformed: "<FUNC> quick brown fox jumps <FUNC> <FUNC> lazy dog"
+
+### Function-Only Variant
+Masks content words with `<CONTENT>` token, preserving only function words (articles, prepositions, conjunctions)
+- **Tests:** Whether models distinguish authors based on grammatical structure
+- **Example transformation:**
+  - Original: "The quick brown fox jumps over the lazy dog"
+  - Transformed: "The <CONTENT> <CONTENT> <CONTENT> <CONTENT> over the <CONTENT> <CONTENT>"
+
+### Part-of-Speech (POS) Variant
+Replaces all words with their POS tags using Universal Dependencies tagset
+- **Tests:** Whether models distinguish authors based on syntactic patterns
+- **Example transformation:**
+  - Original: "The quick brown fox jumps over the lazy dog"
+  - Transformed: "DET ADJ ADJ NOUN VERB ADP DET ADJ NOUN"
+
+### Training Variants
+
+```bash
+# Train a single variant (8 authors × 10 seeds = 80 models per variant)
+./run_llm_stylometry.sh --train --content-only
+./run_llm_stylometry.sh --train --function-only
+./run_llm_stylometry.sh --train --part-of-speech
+
+# Short flags
+./run_llm_stylometry.sh -t -co   # content-only
+./run_llm_stylometry.sh -t -fo   # function-only
+./run_llm_stylometry.sh -t -pos  # part-of-speech
+
+# Train baseline (no variant flag)
+./run_llm_stylometry.sh -t       # baseline (80 models)
+
+# To train all conditions sequentially (baseline + 3 variants = 320 models total):
+./run_llm_stylometry.sh -t                    # baseline
+./run_llm_stylometry.sh -t --content-only     # content variant
+./run_llm_stylometry.sh -t --function-only    # function variant
+./run_llm_stylometry.sh -t --part-of-speech   # POS variant
+```
+
+### Generating Variant Figures
+
+```bash
+# Generate all figures for a single variant
+./run_llm_stylometry.sh --content-only
+./run_llm_stylometry.sh --function-only
+./run_llm_stylometry.sh --part-of-speech
+
+# Generate specific figure for a variant
+./run_llm_stylometry.sh -f 1a --content-only
+./run_llm_stylometry.sh -f 1a --function-only
+
+# Generate baseline figures (no variant flag)
+./run_llm_stylometry.sh           # all baseline figures
+./run_llm_stylometry.sh -f 1a     # specific baseline figure
+
+# To generate all figures for all conditions:
+./run_llm_stylometry.sh                    # baseline
+./run_llm_stylometry.sh --content-only     # content variant
+./run_llm_stylometry.sh --function-only    # function variant
+./run_llm_stylometry.sh --part-of-speech   # POS variant
+```
+
+### Computing Variant Statistics
+
+```bash
+# Single variant statistics
+./run_stats.sh                    # baseline (default)
+./run_stats.sh --content-only     # content variant
+./run_stats.sh --function-only    # function variant
+./run_stats.sh --part-of-speech   # POS variant
+
+# All statistics (baseline + all 3 variants)
+./run_stats.sh --all
+```
+
+### Remote Training with Variants
+
+```bash
+# Train a single variant on GPU server
+./remote_train.sh --content-only
+./remote_train.sh --function-only
+./remote_train.sh --part-of-speech
+
+# Resume variant training
+./remote_train.sh --resume --content-only
+
+# Train baseline on remote server (no variant flag)
+./remote_train.sh
+
+# To train all conditions on remote server, run sequentially:
+./remote_train.sh                    # baseline
+./remote_train.sh --content-only     # content variant
+./remote_train.sh --function-only    # function variant
+./remote_train.sh --part-of-speech   # POS variant
+```
+
+### Model Naming Convention
+
+Models include variant in their directory names:
+- Baseline: `{author}_tokenizer=gpt2_seed={0-9}/`
+- Content: `{author}_variant=content_tokenizer=gpt2_seed={0-9}/`
+- Function: `{author}_variant=function_tokenizer=gpt2_seed={0-9}/`
+- POS: `{author}_variant=pos_tokenizer=gpt2_seed={0-9}/`
+
+### Figure Output Paths
+
+Figures include variant suffix:
+- Baseline: `paper/figs/source/all_losses.pdf`
+- Content: `paper/figs/source/all_losses_content.pdf`
+- Function: `paper/figs/source/all_losses_function.pdf`
+- POS: `paper/figs/source/all_losses_pos.pdf`
+
 ### Using Pre-computed Results
 
 The repository includes pre-computed results from training 80 models (8 authors × 10 random seeds). These results are consolidated in `data/model_results.pkl`.

code/compute_stats.py

@@ -10,10 +10,32 @@
 from pathlib import Path
 from constants import AUTHORS
 
-def load_data():
-    """Load the model results data."""
-    with open('data/model_results.pkl', 'rb') as f:
-        return pickle.load(f)
+def load_data(data_path='data/model_results.pkl', variant=None):
+    """
+    Load and filter model results by variant.
+
+    Args:
+        data_path: Path to consolidated results pickle file
+        variant: One of ['content', 'function', 'pos'] or None for baseline
+
+    Returns:
+        DataFrame filtered to specified variant
+    """
+    with open(data_path, 'rb') as f:
+        df = pickle.load(f)
+
+    # Filter by variant
+    if variant is None:
+        # Baseline: exclude any models with variant column set
+        if 'variant' in df.columns:
+            df = df[df['variant'].isna()].copy()
+    else:
+        # Specific variant: filter to that variant
+        if 'variant' not in df.columns:
+            raise ValueError(f"No variant column in data. Cannot filter for variant '{variant}'")
+        df = df[df['variant'] == variant].copy()
+
+    return df
 
 
 def find_twain_threshold_epoch(df, p_threshold=0.001):
@@ -138,13 +160,32 @@ def generate_author_comparison_table(df):
 
 def main():
     """Main function to compute and display all statistics."""
+    import argparse
+
+    parser = argparse.ArgumentParser(description='Compute statistics for LLM stylometry')
+    parser.add_argument(
+        '--variant',
+        choices=['content', 'function', 'pos'],
+        default=None,
+        help='Analysis variant to compute stats for (default: baseline)'
+    )
+    parser.add_argument(
+        '--data',
+        default='data/model_results.pkl',
+        help='Path to model results file (default: data/model_results.pkl)'
+    )
+
+    args = parser.parse_args()
+
+    # Update header to show variant
+    variant_label = f" (Variant: {args.variant})" if args.variant else " (Baseline)"
     print("=" * 60)
-    print("LLM Stylometry Statistical Analysis")
+    print(f"LLM Stylometry Statistical Analysis{variant_label}")
     print("=" * 60)
 
-    # Load data
+    # Load data with variant filter
     print("\nLoading data...")
-    df = load_data()
+    df = load_data(data_path=args.data, variant=args.variant)
 
     # 1. Find Twain threshold epoch
     print("\n1. Twain Model P-Threshold Analysis")

code/consolidate_model_results.py

@@ -46,15 +46,19 @@ def consolidate_model_results(models_dir='models', output_path='data/model_resul
         dir_name = model_dir.name
         parts = dir_name.split('_')
 
-        # Extract author and seed from directory name
-        # Format: {author}_tokenizer={tokenizer}_seed={seed}
+        # Extract author, variant, tokenizer, and seed from directory name
+        # Baseline format: {author}_tokenizer={tokenizer}_seed={seed}
+        # Variant format: {author}_variant={variant}_tokenizer={tokenizer}_seed={seed}
         author = parts[0]
 
-        # Find tokenizer and seed
+        # Find variant, tokenizer, and seed
+        variant = None
         tokenizer = None
         seed = None
         for part in parts[1:]:
-            if part.startswith('tokenizer='):
+            if part.startswith('variant='):
+                variant = part.split('=')[1]
+            elif part.startswith('tokenizer='):
                 tokenizer = part.split('=')[1]
             elif part.startswith('seed='):
                 seed = int(part.split('=')[1])
@@ -75,6 +79,7 @@ def consolidate_model_results(models_dir='models', output_path='data/model_resul
         # Add model metadata
         df['model_name'] = dir_name
         df['author'] = author
+        df['variant'] = variant  # None for baseline, variant name for variant models
         df['tokenizer'] = tokenizer
         df['checkpoint_path'] = str(model_dir)
 
@@ -103,7 +108,7 @@ def consolidate_model_results(models_dir='models', output_path='data/model_resul
     # Ensure column order matches expected format
     expected_columns = [
         'seed', 'train_author', 'epochs_completed', 'loss_dataset',
-        'loss_value', 'model_name', 'author', 'tokenizer',
+        'loss_value', 'model_name', 'author', 'variant', 'tokenizer',
         'model_config', 'generation_config', 'checkpoint_path'
     ]
 
@@ -128,10 +133,18 @@ def consolidate_model_results(models_dir='models', output_path='data/model_resul
         print(f"Also saved CSV for inspection: {csv_path}")
 
     # Print summary statistics
-    print("\nSummary by author:")
-    summary = consolidated_df.groupby('train_author')['seed'].nunique()
-    for author, num_seeds in summary.items():
-        print(f"  {author}: {num_seeds} seeds")
+    print("\nSummary by author and variant:")
+    if 'variant' in consolidated_df.columns:
+        # Use dropna=False to include None (baseline) values
+        summary = consolidated_df.groupby(['train_author', 'variant'], dropna=False)['seed'].nunique()
+        for (author, variant), num_seeds in summary.items():
+            variant_label = "baseline" if pd.isna(variant) else variant
+            print(f"  {author} ({variant_label}): {num_seeds} seeds")
+    else:
+        # Fallback for old data without variant column
+        summary = consolidated_df.groupby('train_author')['seed'].nunique()
+        for author, num_seeds in summary.items():
+            print(f"  {author}: {num_seeds} seeds")
 
     return consolidated_df
 

code/constants.py

@@ -37,3 +37,29 @@ def find_project_root():
     "fitzgerald",
     "twain",
 ]
+
+# Analysis variants
+ANALYSIS_VARIANTS = ['content', 'function', 'pos']
+
+
+def get_data_dir(variant=None):
+    """
+    Get data directory based on analysis variant.
+
+    Args:
+        variant: One of ANALYSIS_VARIANTS or None for baseline
+
+    Returns:
+        Path to data directory
+    """
+    if variant is None:
+        return CLEANED_DATA_DIR
+
+    if variant not in ANALYSIS_VARIANTS:
+        raise ValueError(f"Invalid variant: {variant}. Must be one of {ANALYSIS_VARIANTS}")
+
+    variant_dir = CLEANED_DATA_DIR / f"{variant}_only"
+    if not variant_dir.exists():
+        raise FileNotFoundError(f"Variant directory not found: {variant_dir}")
+
+    return variant_dir