File size: 7,239 Bytes
6aef776
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
import os
import json
import pickle
import numpy as np
from tqdm import tqdm
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, compute_tsne_with_cache, compute_precomputed_regions
from utils.ui import update_task_display, instance_to_df
from utils.interp_space_utils import cached_generate_style_embedding, compute_g2v_features, compute_predicted_author

def load_config(path="config/config.yaml"):
    with open(path, "r") as f:
        return yaml.safe_load(f)

def precompute_all_caches(
    models_to_test=None,
    instances_to_process=None,
    config_path="config/config.yaml",
    force_regenerate=False
):
    """
    Precompute all cache files using the EXACT same methods as app.py.
    This follows the exact flow: load_task β†’ update_task_display β†’ run_visualization
    """
    
    if models_to_test is None:
        models_to_test = [
            'gabrielloiseau/LUAR-MUD-sentence-transformers',
            'gabrielloiseau/LUAR-CRUD-sentence-transformers', 
            'miladalsh/light-luar',
            'AnnaWegmann/Style-Embedding'
        ]
    
    print("=" * 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']
    
    if instances_to_process is None:
        instances_to_process = instance_ids
    
    print(f"Processing {len(instances_to_process)} instances with {len(models_to_test)} models")
    
    total_combinations = len(models_to_test) * len(instances_to_process)
    current_combination = 0
    
    cache_stats = {
        'embeddings_generated': 0,
        'tsne_computed': 0,
        'regions_computed': 0,
        'errors': []
    }
    
    for model_name in models_to_test:
        print(f"\n{'=' * 40}")
        print(f"PROCESSING MODEL: {model_name}")
        print(f"{'=' * 40}")
        
        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}")
                
                # STEP 1: Replicate the exact flow from load_button.click()
                print("  β†’ Replicating load_button.click() flow...")
                
                # Create ground truth (using placeholder since we're caching)
                ground_truth_author = None  # Will be determined by the instance data
                
                # Call update_task_display EXACTLY like app.py does
                task_results = update_task_display(
                    mode="Predefined HRS Task",  # Always use predefined for caching
                    iid=f"Task {instance_id}",
                    instances=instances,
                    background_df=clustered_authors_df,
                    mystery_file=None,  # Not used for predefined
                    cand1_file=None,   # Not used for predefined  
                    cand2_file=None,   # Not used for predefined
                    cand3_file=None,   # Not used for predefined
                    true_author=ground_truth_author,
                    model_radio=model_name,
                    custom_model_input=""
                )
                
                # Extract the results exactly like app.py expects
                (header_html, mystery_html, c0_html, c1_html, c2_html, 
                 mystery_state, c0_state, c1_state, c2_state, 
                 task_authors_embeddings_df, background_authors_embeddings_df, 
                 predicted_author, ground_truth_author) = task_results
                
                print(f"    βœ“ Embeddings generated for {len(task_authors_embeddings_df)} task authors")
                print(f"    βœ“ Background embeddings: {len(background_authors_embeddings_df)} authors")
                cache_stats['embeddings_generated'] += 1
                
                # STEP 2: Replicate the exact flow from run_btn.click()
                print("  β†’ Replicating run_btn.click() flow...")
                
                # Call visualize_clusters_plotly EXACTLY like app.py does
                viz_results = visualize_clusters_plotly(
                    iid=int(instance_id),
                    cfg=cfg,
                    instances=instances,
                    model_radio=model_name,
                    custom_model_input="",
                    task_authors_df=task_authors_embeddings_df,
                    background_authors_embeddings_df=background_authors_embeddings_df,
                    pred_idx=predicted_author,
                    gt_idx=ground_truth_author
                )
                
                # Extract results exactly like app.py expects
                (fig, style_names, bg_proj, bg_ids, bg_authors_df, 
                 precomputed_regions_state, precomputed_regions_radio) = viz_results
                
                print(f"    βœ“ t-SNE projection computed")
                print(f"    βœ“ Precomputed regions generated")
                cache_stats['tsne_computed'] += 1
                cache_stats['regions_computed'] += 1
                
                print(f"  βœ“ Instance {instance_id} with model {model_name} completed successfully")
                
            except Exception as e:
                error_msg = f"Error processing instance {instance_id} with model {model_name}: {str(e)}"
                print(f"  βœ— {error_msg}")
                cache_stats['errors'].append(error_msg)
                import traceback
                traceback.print_exc()
                continue
    
    # Print final statistics
    print("\n" + "=" * 60)
    print("CACHE PRECOMPUTATION COMPLETED")
    print("=" * 60)
    print(f"Embeddings generated: {cache_stats['embeddings_generated']}")
    print(f"t-SNE projections computed: {cache_stats['tsne_computed']}")
    print(f"Region sets computed: {cache_stats['regions_computed']}")
    print(f"Errors encountered: {len(cache_stats['errors'])}")
    
    if cache_stats['errors']:
        print("\nERROR DETAILS:")
        for error in cache_stats['errors']:
            print(f"  - {error}")
    
    return cache_stats

# Import the exact functions your app uses
from utils.visualizations import visualize_clusters_plotly

if __name__ == "__main__":
    # Test with a small subset first
    instances=[i for i in range(2)]  # First 2 instances for testing
    cache_stats = precompute_all_caches(
        models_to_test=[
            'gabrielloiseau/LUAR-MUD-sentence-transformers'
        ],
        instances_to_process=instances,  
        force_regenerate=False
    )
    
    print(f"\nCache precomputation completed with {len(cache_stats['errors'])} errors.")