<jupyter_start><jupyter_text>Baseline model using a pre-trained VGGish model<jupyter_code>class WAV2VEC(nn.Module):
    def __init__(self, input_size=768, hidden_size=64):
        super().__init__()

        self.wav2vec2 = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base-960h")

        self.adaptor = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(True),           
            nn.Dropout(0.1),            
            nn.Linear(hidden_size, input_size),
        )  
        
        self.classifier = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(True),           
            nn.Dropout(0.1),            
            nn.Linear(hidden_size, 1),
        )           

    def forward(self, x):
        # x shape: (B,E)
        x = self.wav2vec2(x).last_hidden_state
        
        x = self.adaptor(x)
        
        # pooling
        x, _ = x.max(dim=1)

        # Mutilayer perceptron
        out = self.classifier(x)
        # out shape: (B,1)

        # Remove last dimension
        return out.squeeze(-1)
        # return shape: (B)
def train(loader, model, criterion, optimizer, epoch, cuda, log_interval, max_norm=1, verbose=True):
    model.train()
    global_epoch_loss = 0
    samples = 0
    for batch_idx, (_, data, target) in enumerate(loader):
        if cuda:
            data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target.float())
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), max_norm)
        optimizer.step()
        global_epoch_loss += loss.data.item() * len(target)
        samples += len(target)
        if verbose and (batch_idx % log_interval == 0):
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, samples, len(loader.dataset), 100*samples/len(loader.dataset), global_epoch_loss/samples))
    return global_epoch_loss / samples
def test(loader, model, criterion, cuda, verbose=True, data_set='Test', save=None):
    model.eval()
    test_loss = 0
    tpred = []
    ttarget = []

    if save is not None:
        csv = open(save, 'wt')
        print('index,prob', file=csv)

    with torch.no_grad():
        for keys, data, target in loader:
            if cuda:
                data, target = data.cuda(), target.cuda()
            output = model(data)
            pred = output.sigmoid()
            tpred.append(pred.cpu().numpy())

            if target[0] != -1:
                loss = criterion(output, target.float()).data.item()
                test_loss += loss * len(target) # sum up batch loss 
                ttarget.append(target.cpu().numpy())

            if save is not None:
                for i, key in enumerate(keys):
                    print(f'{key},{pred[i]}', file=csv)
    
    if len(ttarget) > 0:
        test_loss /= len(loader.dataset)
        auc, auc_ci = roc_auc_score_ci(np.concatenate(ttarget), np.concatenate(tpred))
        if verbose:
            print('\n{} set: Average loss: {:.4f}, AUC: {:.1f}% ({:.1f}% - {:.1f}%)\n'.format(
                data_set, test_loss, 100 * auc, auc_ci[0]*100, auc_ci[1]*100))

        return test_loss, auc
args = SimpleNamespace(
    # general options
    train_path = '../input/covid/train',         # train data folder
    valid_path = '../input/covid/valid',         # valid data folder
    test_path = '../input/covid/test',           # test data folder
    batch_size = 16,                             # training and valid batch size
    test_batch_size = 16,                        # batch size for testing
    arch = 'WAV2VEC',                            # WAV2VEC
    epochs = 50,                                 # maximum number of epochs to train
    lr = 0.0002,                                 # learning rate
    momentum = 0.9,                              # SGD momentum, for SGD only
    optimizer = 'adam',                          # optimization method: sgd | adam
    seed = seed,                                 # random seed
    log_interval = 5,                            # how many batches to wait before logging training status
    patience = 5,                                # how many epochs of no loss improvement should we wait before stop training
    checkpoint = '.',                            # checkpoints directory
    train = True,                                # train before testing
    cuda = True,                                 # use gpu
    num_workers = 2,                             # how many subprocesses to use for data loading
    hidden_size = 64
)
args.cuda = args.cuda and torch.cuda.is_available()
if args.cuda:
    print('Using CUDA with {0} GPUs'.format(torch.cuda.device_count()))

# build model
model = WAV2VEC(hidden_size=args.hidden_size)
if args.cuda:
    model.cuda()

# Define criterion
criterion = nn.BCEWithLogitsLoss(reduction='mean') # This loss combines a Sigmoid layer and the BCELoss in one single class.<jupyter_output>Using CUDA with 1 GPUs<jupyter_text>Train model<jupyter_code># partial freeze of wac2vec parameters. Only feature_projection parameters are fine tuned
for param in model.wav2vec2.feature_extractor.parameters():
    param.requires_grad = False
for param in model.wav2vec2.encoder.parameters():
    param.requires_grad = False
model.wav2vec2
# loading data
if args.train:
    train_dataset = Loader(args.train_path)
    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers)

    valid_dataset = Loader(args.valid_path)
    valid_loader = torch.utils.data.DataLoader(
        valid_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)

    # define optimizer
    if args.optimizer.lower() == 'adam':
        optimizer = optim.Adam(model.parameters(), lr=args.lr)
    else:
        optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

    best_valid_auc = 0
    iteration = 0
    epoch = 1
    best_epoch = epoch
    
    # trainint with early stopping
    t0 = time.time()
    while (epoch < args.epochs + 1) and (iteration < args.patience):
        train(train_loader, model, criterion, optimizer, epoch, args.cuda, args.log_interval)
        valid_loss, valid_auc = test(valid_loader, model, criterion, args.cuda, data_set='Validation')
        if not os.path.isdir(args.checkpoint):
            os.mkdir(args.checkpoint)
        torch.save(model.state_dict(), './{}/model{:03d}.pt'.format(args.checkpoint, epoch))
        if valid_auc <= best_valid_auc:
            iteration += 1
            print('AUC was not improved, iteration {0}'.format(str(iteration)))
        else:
            print('Saving state')
            iteration = 0
            best_valid_auc = valid_auc
            best_epoch = epoch
            state = {
                'valid_auc': valid_auc,
                'valid_loss': valid_loss,
                'epoch': epoch,
            }
            if not os.path.isdir(args.checkpoint):
                os.mkdir(args.checkpoint)
            torch.save(state, './{}/ckpt.pt'.format(args.checkpoint))
        epoch += 1
        print(f'Elapsed seconds: ({time.time() - t0:.0f}s)')
    print(f'Best AUC: {best_valid_auc*100:.1f}% on epoch {best_epoch}')<jupyter_output>Train Epoch: 1 [16/2160 (1%)]	Loss: 0.712866
Train Epoch: 1 [96/2160 (4%)]	Loss: 0.699255
Train Epoch: 1 [176/2160 (8%)]	Loss: 0.709024
Train Epoch: 1 [256/2160 (12%)]	Loss: 0.706287
Train Epoch: 1 [336/2160 (16%)]	Loss: 0.703815
Train Epoch: 1 [416/2160 (19%)]	Loss: 0.702599
Train Epoch: 1 [496/2160 (23%)]	Loss: 0.700754
Train Epoch: 1 [576/2160 (27%)]	Loss: 0.699665
Train Epoch: 1 [656/2160 (30%)]	Loss: 0.698664
Train Epoch: 1 [736/2160 (34%)]	Loss: 0.697913
Train Epoch: 1 [816/2160 (38%)]	Loss: 0.697749
Train Epoch: 1 [896/2160 (41%)]	Loss: 0.697388
Train Epoch: 1 [976/2160 (45%)]	Loss: 0.697199
Train Epoch: 1 [1056/2160 (49%)]	Loss: 0.696284
Train Epoch: 1 [1136/2160 (53%)]	Loss: 0.696289
Train Epoch: 1 [1216/2160 (56%)]	Loss: 0.696081
Train Epoch: 1 [1296/2160 (60%)]	Loss: 0.695670
Train Epoch: 1 [1376/2160 (64%)]	Loss: 0.695931
Train Epoch: 1 [1456/2160 (67%)]	Loss: 0.696007
Train Epoch: 1 [1536/2160 (71%)]	Loss: 0.695941
Train Epoch: 1 [1616/2160 (75%)]	Loss: 0.695878
Train Epoc[...]<jupyter_text>Test Model<jupyter_code>test_dataset = Loader(args.test_path)
test_loader = torch.utils.data.DataLoader(
    test_dataset, batch_size=args.test_batch_size, shuffle=False, num_workers=args.num_workers)

# get best epoch and model
state = torch.load('./{}/ckpt.pt'.format(args.checkpoint))
epoch = state['epoch']
print("Testing model (epoch {})".format(epoch))
model.load_state_dict(torch.load('./{}/model{:03d}.pt'.format(args.checkpoint, epoch)))
if args.cuda:
    model.cuda()

results = 'submission.csv'
print("Saving results in {}".format(results))
test(test_loader, model, criterion, args.cuda, save=results)<jupyter_output>Testing model (epoch 17)
Saving results in submission.csv
{'fraction_non_alphanumeric': 0.12623925974884337, 'fraction_numerical': 0.05111257986340604, 'mean_word_length': 2.303857350800582, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '94925', 'license': 'MIT', 'n_tokens_mistral': 3336, 'n_tokens_neox': 2955, 'n_words': 706, 'path': 'examples/f4-wav2vec2.ipynb', 'repo_name': 'shubham-kulkarniS/jaafe', 'size': '9078'}