File size: 691 Bytes
8a0e04e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import nbformat from 'nbformat';
import { parse } from 'ast';

export class NotebookParser {
  constructor() {
    this.TARGET_PARAMS = new Set([
      'learning_rate', 'lr', 'batch_size', 'epochs', 'num_epochs',
      'optimizer', 'loss', 'dropout', 'weight_decay', 'momentum',
      'hidden_dim', 'num_layers', 'activation', 'scheduler',
      'backbone', 'model_name', 'pretrained', 'num_classes',
      'augmentation', 'early_stopping', 'gradient_clip', 'warmup_epochs'
    ]);
    this.extractedParams = [];
  }

  async parseNotebook(notebookPath) {
    const content = await fs.readFile(notebookPath, 'utf-8');
    const notebook = JSON.parse(content);

    this.extractedParams = [];