OTExplorer/example.py at main · Jank14/OTExplorer · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
"""
Simple example demonstrating the OT analysis pipeline.
"""

import torch
import numpy as np
from transformer_utils import ActivationExtractor
from ot_analysis import OTAnalyzer
from visualization import OTVisualizer

def simple_example():
    """Run a simple example with a few sample texts."""

    print("Simple OT Analysis Example")
    print("=" * 40)

    # Sample texts
    sample_texts = [
        "The quick brown fox jumps over the lazy dog",
        "Machine learning is a subset of artificial intelligence",
        "Optimal transport measures the cost of moving mass",
        "Transformers have revolutionized natural language processing",
        "Emergence occurs when complex behaviors arise from simple rules"
    ]

    print(f"Using {len(sample_texts)} sample texts")

    # Initialize components
    print("Loading GPT-2 model...")
    extractor = ActivationExtractor(model_name="gpt2", device="cpu")

    print("Extracting activations...")
    activations = extractor.extract_activations_batch(sample_texts)

    print("Preparing representations...")
    # Get all layer representations
    all_reps = extractor.get_all_layer_representations(activations)

    # Convert to numpy and prepare for analysis
    layer_representations = []
    for i, reps in enumerate(all_reps):
        reps_np = reps.numpy()
        batch_size, seq_len, hidden_dim = reps_np.shape

        # Reshape to combine batch and sequence dimensions
        reps_reshaped = reps_np.reshape(-1, hidden_dim)

        # Remove padding tokens (simple heuristic)
        non_padding_mask = np.any(reps_reshaped != 0, axis=1)
        reps_clean = reps_reshaped[non_padding_mask]

        layer_representations.append(reps_clean)
        print(f"Layer {i}: {reps_clean.shape}")

    # Run OT analysis
    print("Running OT analysis...")
    ot_analyzer = OTAnalyzer(method="emd", normalize=True)
    analysis_results = ot_analyzer.analyze_representation_flow(layer_representations)

    # Create visualizations
    print("Creating visualizations...")
    visualizer = OTVisualizer()

    # Plot OT distances
    visualizer.plot_ot_distances(analysis_results['ot_distances'])

    # Plot phase transitions
    visualizer.plot_phase_transitions(
        analysis_results['ot_distances'],
        analysis_results['phase_transitions']
    )

    # Plot entropy evolution
    visualizer.plot_entropy_evolution(analysis_results['entropies'])

    # Print summary
    print("\nResults Summary:")
    print(f"Number of layers: {analysis_results['num_layers']}")
    print(f"OT distances: {[f'{d:.4f}' for d in analysis_results['ot_distances']]}")
    print(f"Phase transitions: {analysis_results['phase_transitions']}")
    print(f"Entropies: {[f'{e:.4f}' for e in analysis_results['entropies']]}")

    # Cleanup
    extractor.cleanup()

    print("\nExample completed!")

def compare_methods_example():
    """Compare different OT methods."""

    print("\nComparing OT Methods")
    print("=" * 40)

    # Sample texts
    sample_texts = [
        "Hello world",
        "Python programming",
        "Data science",
        "Artificial intelligence"
    ]

    # Initialize components
    extractor = ActivationExtractor(model_name="gpt2", device="cpu")
    activations = extractor.extract_activations_batch(sample_texts)

    # Prepare representations
    all_reps = extractor.get_all_layer_representations(activations)
    layer_representations = []

    for reps in all_reps:
        reps_np = reps.numpy()
        batch_size, seq_len, hidden_dim = reps_np.shape
        reps_reshaped = reps_np.reshape(-1, hidden_dim)
        non_padding_mask = np.any(reps_reshaped != 0, axis=1)
        reps_clean = reps_reshaped[non_padding_mask]
        layer_representations.append(reps_clean)

    # Compare methods
    methods = ["emd", "sinkhorn"]
    results = {}

    for method in methods:
        print(f"Testing {method} method...")
        ot_analyzer = OTAnalyzer(method=method, normalize=True)
        distances = ot_analyzer.compute_layer_distances(layer_representations)
        results[method] = distances

    # Print comparison
    print("\nMethod Comparison:")
    for i in range(len(results["emd"])):
        print(f"L{i}-L{i+1}: EMD={results['emd'][i]:.4f}, Sinkhorn={results['sinkhorn'][i]:.4f}")

    # Cleanup
    extractor.cleanup()

if __name__ == "__main__":
    # Run simple example
    simple_example()

    # Run method comparison
    compare_methods_example()