SSL-NL-eval/phone_LDA.py at main · mdhk/SSL-NL-eval · 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
"""
Example usage:
python phone_LDA.py --embeddings_file="embeddings/amsterdamNLP_Wav2Vec2-NL_phone_embs.pkl" --model_name="w2v2-nl" --subset="MLS"
"""

import pickle
import pandas as pd
from tqdm import tqdm
from scipy.stats import zscore
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.metrics import silhouette_samples
from argparse import ArgumentParser
from pathlib import Path

from model_utils import model_depth_map

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "--embeddings_file",
        required=True,
        type=str,
        help='filepath to the pkl file with extracted phone embeddings',
    )
    parser.add_argument(
        "--model_name",
        type=str,
        required=True,
        help='name of model that generated the embeddings (for file saving)'
    )
    parser.add_argument(
        "--annotations_file",
        type=str,
        default='SSL-NL/annotations/phone_annotations.csv',
        help='filepath to the phone annotations file'
    )
    parser.add_argument(
        "--subset",
        type=str,
        default='all',
        help='subset to evaluate embeddings on (MLS or IFADV)'
    )
    parser.add_argument(
        "--results_dir",
        type=str,
        default='results',
        help='directory to store results to'
    )
    args, unk_args = parser.parse_known_args()

    embeddings = pickle.load(open(args.embeddings_file, 'rb'))
    annotations = pd.read_csv(args.annotations_file)
    results_dir = Path(args.results_dir)
    results_dir.mkdir(exist_ok=True, parents=True)

    N_layers = len(embeddings.keys())

    lda_item_scores = []
    lda_projections = {
        layer: None
        for layer in embeddings.keys()
    }

    print(f'Computing LDA clustering scores for {args.model_name} on {args.subset} data...')

    for layer in tqdm(embeddings.keys(), desc='\tclustering & scoring by layer'):
        embs = zscore(embeddings[layer])

        if args.subset != 'all':
            train_idx = (annotations['subset'] == args.subset) & (annotations['split'] == 'train')
            test_idx = (annotations['subset'] == args.subset) & (annotations['split'] == 'test')
        else:
            train_idx = annotations['split'] == 'train'
            test_idx = annotations['split'] == 'test'

        X_train = embs[train_idx]
        X_test = embs[test_idx]
        y_train = annotations['phone'][train_idx].values
        y_test = annotations['phone'][test_idx].values

        n_comps = len(annotations['phone'].unique())-1
        lda = LinearDiscriminantAnalysis(n_components=n_comps)
        lda.fit(X_train, y_train)

        lda_projections[layer] = lda.transform(embs[test_idx])

        silh_sample_scores = silhouette_samples(
            lda_projections[layer],
            y_test,
            metric='cosine'
        )

        for i in range(len(silh_sample_scores)):
            lda_item_scores.append((args.subset, args.model_name, layer, model_depth_map[layer],
                                    y_test[i], 'silhouette_score', silh_sample_scores[i]))

    lda_score_df = pd.DataFrame(lda_item_scores,
                                columns=['eval_set', 'model_name', 'layer', 'model_depth',
                                         'target', 'score_name', 'score'])

    lda_score_df.to_csv(results_dir / f'phone_LDA_item-scores_{args.subset}_{args.model_name}.csv', index=False)

    print(f'Done! Saved results to {args.results_dir}')