File size: 9,641 Bytes
29792f4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


"""
To run this script, from the root of the repo. Make sure to have Flask installed

    FLASK_DEBUG=1 FLASK_APP=scripts.mos flask run -p 4567
    # or if you have gunicorn
    gunicorn -w 4 -b 127.0.0.1:8895 -t 120 'scripts.mos:app'  --access-logfile -

"""
from collections import defaultdict
from functools import wraps
from hashlib import sha1
import json
import math
from pathlib import Path
import random
import typing as tp

from flask import Flask, redirect, render_template, request, session, url_for

from audiocraft import train
from audiocraft.utils.samples.manager import get_samples_for_xps


SAMPLES_PER_PAGE = 8
MAX_RATING = 5
storage = Path(train.main.dora.dir / 'mos_storage')
storage.mkdir(exist_ok=True)
surveys = storage / 'surveys'
surveys.mkdir(exist_ok=True)
magma_root = Path(train.__file__).parent.parent
app = Flask('mos', static_folder=str(magma_root / 'scripts/static'),
            template_folder=str(magma_root / 'scripts/templates'))
app.secret_key = b'audiocraft makes the best songs'


def normalize_path(path: Path):
    """Just to make path a bit nicer, make them relative to the Dora root dir.
    """
    path = path.resolve()
    dora_dir = train.main.dora.dir.resolve() / 'xps'
    return path.relative_to(dora_dir)


def get_full_path(normalized_path: Path):
    """Revert `normalize_path`.
    """
    return train.main.dora.dir.resolve() / 'xps' / normalized_path


def get_signature(xps: tp.List[str]):
    """Return a signature for a list of XP signatures.
    """
    return sha1(json.dumps(xps).encode()).hexdigest()[:10]


def ensure_logged(func):
    """Ensure user is logged in.
    """
    @wraps(func)
    def _wrapped(*args, **kwargs):
        user = session.get('user')
        if user is None:
            return redirect(url_for('login', redirect_to=request.url))
        return func(*args, **kwargs)
    return _wrapped


@app.route('/login', methods=['GET', 'POST'])
def login():
    """Login user if not already, then redirect.
    """
    user = session.get('user')
    if user is None:
        error = None
        if request.method == 'POST':
            user = request.form['user']
            if not user:
                error = 'User cannot be empty'
        if user is None or error:
            return render_template('login.html', error=error)
    assert user
    session['user'] = user
    redirect_to = request.args.get('redirect_to')
    if redirect_to is None:
        redirect_to = url_for('index')
    return redirect(redirect_to)


@app.route('/', methods=['GET', 'POST'])
@ensure_logged
def index():
    """Offer to create a new study.
    """
    errors = []
    if request.method == 'POST':
        xps_or_grids = [part.strip() for part in request.form['xps'].split()]
        xps = set()
        for xp_or_grid in xps_or_grids:
            xp_path = train.main.dora.dir / 'xps' / xp_or_grid
            if xp_path.exists():
                xps.add(xp_or_grid)
                continue
            grid_path = train.main.dora.dir / 'grids' / xp_or_grid
            if grid_path.exists():
                for child in grid_path.iterdir():
                    if child.is_symlink():
                        xps.add(child.name)
                continue
            errors.append(f'{xp_or_grid} is neither an XP nor a grid!')
        assert xps or errors
        blind = 'true' if request.form.get('blind') == 'on' else 'false'
        xps = list(xps)
        if not errors:
            signature = get_signature(xps)
            manifest = {
                'xps': xps,
            }
            survey_path = surveys / signature
            survey_path.mkdir(exist_ok=True)
            with open(survey_path / 'manifest.json', 'w') as f:
                json.dump(manifest, f, indent=2)
            return redirect(url_for('survey', blind=blind, signature=signature))
    return render_template('index.html', errors=errors)


@app.route('/survey/<signature>', methods=['GET', 'POST'])
@ensure_logged
def survey(signature):
    success = request.args.get('success', False)
    seed = int(request.args.get('seed', 4321))
    blind = request.args.get('blind', 'false') in ['true', 'on', 'True']
    exclude_prompted = request.args.get('exclude_prompted', 'false') in ['true', 'on', 'True']
    exclude_unprompted = request.args.get('exclude_unprompted', 'false') in ['true', 'on', 'True']
    max_epoch = int(request.args.get('max_epoch', '-1'))
    survey_path = surveys / signature
    assert survey_path.exists(), survey_path

    user = session['user']
    result_folder = survey_path / 'results'
    result_folder.mkdir(exist_ok=True)
    result_file = result_folder / f'{user}_{seed}.json'

    with open(survey_path / 'manifest.json') as f:
        manifest = json.load(f)

    xps = [train.main.get_xp_from_sig(xp) for xp in manifest['xps']]
    names, ref_name = train.main.get_names(xps)

    samples_kwargs = {
        'exclude_prompted': exclude_prompted,
        'exclude_unprompted': exclude_unprompted,
        'max_epoch': max_epoch,
    }
    matched_samples = get_samples_for_xps(xps, epoch=-1, **samples_kwargs)  # fetch latest epoch
    models_by_id = {
        id: [{
            'xp': xps[idx],
            'xp_name': names[idx],
            'model_id': f'{xps[idx].sig}-{sample.id}',
            'sample': sample,
            'is_prompted': sample.prompt is not None,
            'errors': [],
        } for idx, sample in enumerate(samples)]
        for id, samples in matched_samples.items()
    }
    experiments = [
        {'xp': xp, 'name': names[idx], 'epoch': list(matched_samples.values())[0][idx].epoch}
        for idx, xp in enumerate(xps)
    ]

    keys = list(matched_samples.keys())
    keys.sort()
    rng = random.Random(seed)
    rng.shuffle(keys)
    model_ids = keys[:SAMPLES_PER_PAGE]

    if blind:
        for key in model_ids:
            rng.shuffle(models_by_id[key])

    ok = True
    if request.method == 'POST':
        all_samples_results = []
        for id in model_ids:
            models = models_by_id[id]
            result = {
                'id': id,
                'is_prompted': models[0]['is_prompted'],
                'models': {}
            }
            all_samples_results.append(result)
            for model in models:
                rating = request.form[model['model_id']]
                if rating:
                    rating = int(rating)
                    assert rating <= MAX_RATING and rating >= 1
                    result['models'][model['xp'].sig] = rating
                    model['rating'] = rating
                else:
                    ok = False
                    model['errors'].append('Please rate this model.')
        if ok:
            result = {
                'results': all_samples_results,
                'seed': seed,
                'user': user,
                'blind': blind,
                'exclude_prompted': exclude_prompted,
                'exclude_unprompted': exclude_unprompted,
            }
            print(result)
            with open(result_file, 'w') as f:
                json.dump(result, f)
            seed = seed + 1
            return redirect(url_for(
                'survey', signature=signature, blind=blind, seed=seed,
                exclude_prompted=exclude_prompted, exclude_unprompted=exclude_unprompted,
                max_epoch=max_epoch, success=True))

    ratings = list(range(1, MAX_RATING + 1))
    return render_template(
        'survey.html', ratings=ratings, blind=blind, seed=seed, signature=signature, success=success,
        exclude_prompted=exclude_prompted, exclude_unprompted=exclude_unprompted, max_epoch=max_epoch,
        experiments=experiments, models_by_id=models_by_id, model_ids=model_ids, errors=[],
        ref_name=ref_name, already_filled=result_file.exists())


@app.route('/audio/<path:path>')
def audio(path: str):
    full_path = Path('/') / path
    assert full_path.suffix in [".mp3", ".wav"]
    return full_path.read_bytes(), {'Content-Type': 'audio/mpeg'}


def mean(x):
    return sum(x) / len(x)


def std(x):
    m = mean(x)
    return math.sqrt(sum((i - m)**2 for i in x) / len(x))


@app.route('/results/<signature>')
@ensure_logged
def results(signature):

    survey_path = surveys / signature
    assert survey_path.exists(), survey_path
    result_folder = survey_path / 'results'
    result_folder.mkdir(exist_ok=True)

    # ratings per model, then per user.
    ratings_per_model = defaultdict(list)
    users = []
    for result_file in result_folder.iterdir():
        if result_file.suffix != '.json':
            continue
        with open(result_file) as f:
            results = json.load(f)
        users.append(results['user'])
        for result in results['results']:
            for sig, rating in result['models'].items():
                ratings_per_model[sig].append(rating)

    fmt = '{:.2f}'
    models = []
    for model in sorted(ratings_per_model.keys()):
        ratings = ratings_per_model[model]

        models.append({
            'sig': model,
            'samples': len(ratings),
            'mean_rating': fmt.format(mean(ratings)),
            # the value 1.96 was probably chosen to achieve some
            # confidence interval assuming gaussianity.
            'std_rating': fmt.format(1.96 * std(ratings) / len(ratings)**0.5),
        })
    return render_template('results.html', signature=signature, models=models, users=users)