# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="bIioeTx3pgIG"
# #Requirements

# + colab={"base_uri": "https://localhost:8080/"} id="1UTOca7S_2zj" outputId="43f0802c-6479-4b62-e8c7-b0ad69e070d0"
# !ls

# + colab={"base_uri": "https://localhost:8080/"} id="AiyNkyuSmHKc" outputId="dba09cbe-a4ac-48f9-e71e-838adfb3a83a"
# !nvidia-smi

# + colab={"base_uri": "https://localhost:8080/"} id="S2gkMaBbabVB" outputId="be4441d0-3c10-43a7-873f-721213495680"
from google.colab import drive
drive.mount('/content/drive')

# + [markdown] id="kaYmOTOJs3hU"
# ## Scripts

# + colab={"base_uri": "https://localhost:8080/"} id="r0ueAmNUt8st" outputId="ef276c5a-9bd9-4d18-ce84-0499cde362c3"
# %%writefile init.sh

pip install soundfile catalyst -q
pip install torchlibrosa -q

# mkdir data/ data/raw/

unzip -q data/raw/audio_files.zip -d data/
unzip -q data/raw/AdditionalUtterances.zip -d data/
unzip -q data/raw/nlp_keywords_29Oct2020.zip -d data/

# + colab={"base_uri": "https://localhost:8080/"} id="rKlvcZYfskqH" outputId="7b867f82-dd5d-45be-ccdc-f9fc386de901"
# %%writefile init.py

import os, sys, gc, glob
import argparse
import pandas as pd
import numpy as np


parser = argparse.ArgumentParser()

parser.add_argument('-prefix', type=str, default='./data/here/', help="postprocess folder")
parser.add_argument('-data', type=str, default='data/raw/', help="where to fing the zipfiles")


def preprocessing(args):
	def what(splits, ns=-3):
		return '_'.join([ splits[ns], splits[-1][:-4] ]).lower()
	
	additionnal = glob.glob("data/latest_keywords/*/*.wav")
	additionnal += glob.glob("data/nlp_keywords/*/*.wav")

	os.makedirs(args.prefix, exist_ok=True)

	add = pd.DataFrame({'fn': additionnal})
	add['bname'] = add['fn'].apply(lambda x: what(x.split('/')))
	add['type'] = 'add'
	add['target'] = add['fn'].apply(lambda x: x.split('/')[-2])
	add.to_csv(args.prefix + 'AddTrain.csv', index=False)

	train = pd.read_csv(args.data + 'Train.csv')
	train['bname'] = train['fn'].apply(lambda x: what(x.split('/'), -2))
	train['type'] = 'base'
	train['fn'] = 'data/' + train['fn']
	train.rename(columns = {'label': 'target'}, inplace=True)
	train.to_csv(args.prefix + 'BaseTrain.csv', index=False)

	train = pd.concat([train, add], axis=0)
	train.to_csv(args.prefix + 'Train.csv', index=False)

	subs = pd.read_csv(args.data + 'SampleSubmission.csv')
	subs['fn'] = 'data/' + subs['fn']
	subs['bname'] = subs['fn'].apply(lambda x: what(x.split('/'), -2))
	cols = subs.columns.tolist()
	subs = subs[[cols[0]] + [cols[-1]] + cols[1:-1]]
	subs.to_csv(args.prefix + 'SampleSubmission.csv', index=False)


if __name__ == '__main__':
	args = parser.parse_args()
	preprocessing(args)

# + [markdown] id="QfkGIyjgupyD"
# ## Data Download

# + id="9E3bsRoOuvkj"
# !chmod +x init.sh
# !./init.sh

# + id="MuJUwcczzTZd"
# !python init.py

# + [markdown] id="a9X2-Q-5WTkM"
# #Imports

# + id="8XObdABGhfVa"
import warnings
warnings.simplefilter('ignore')

# + id="YSARNegOWWZd"
import os
import gc
import sys
import h5py
import cv2
import glob
import math
import random
import librosa
import zipfile
import numpy as np
import pandas as pd
from librosa import display as libdisplay
from tqdm.notebook import tqdm

# + id="-1b7hm-2OtAu"
import torchlibrosa
from torchlibrosa.stft import Spectrogram, LogmelFilterBank
from torchlibrosa.augmentation import SpecAugmentation

# + id="7TCDPB0mMTnd"
from torchvision import models

# + id="cRna23dmG4Uw"
from sklearn.model_selection import train_test_split, StratifiedKFold
from sklearn.metrics import log_loss

# + id="BWnNKykv3kmz"
from catalyst.contrib.nn.criterion import FocalLossMultiClass

# + id="lOdN59az6Mhj"
import torch
import torch.nn as nn
import torch.nn.functional as F
from keras.utils import to_categorical

# + id="HuhiKmCtchG-"
import IPython.display as ipd
from matplotlib import pyplot as plt

# + [markdown] id="3auEtamPiFL9"
# #Envs

# + id="6Uxu_L7dayhc"
path = 'data/'

# + id="x46mDO1tiGaV"
seed = 1999

# + id="kaHzC4X5_uXv"
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)

if torch.cuda.is_available(): 
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# + id="u-94NUtOEqVz"
os.makedirs('MODELS/', exist_ok=True)

# + id="wxwzxxjYV1XQ"
# #Placeholder for the training and test spectogram's images
# #It is going to store the spec, we will shortly generate.
# os.makedirs('Imgs/Train/', exist_ok=True)
# os.makedirs('Imgs/Test/', exist_ok=True)

# + [markdown] id="4TERaIqw5_lA"
# #Utilities

# + [markdown] id="DrAS-iagYNde"
# ## Basic functions

# + id="_g5DLzdQlNMn"
from joblib import Parallel, delayed
import multiprocessing
import time


# + id="j8bQqkjUm6Ow"
def pad_or_truncate(x, audio_length):
  """Pad all audio to specific length."""
  if len(x) <= audio_length:
      return np.concatenate((x, np.zeros(audio_length - len(x))), axis=0)
  return x[:audio_length]

def load_hdf5(hdf5_path):
  hf = h5py.File(hdf5_path, 'r')
  audio_name = hf['audio_name'][:].tolist()
  waveform = hf['waveform']
  target = hf['target'][:].tolist()
  return audio_name, waveform, target, hf

def load_npy(npy_path):
  return np.load(npy_path)

def pack_waveforms_to_npy(npy_path, df, sr=44100, secs=3):
  """Pack waveform and target of several audio clips to a single hdf5 file. 
  This can speed up loading and training.
  """
  def __parallel(df, w, n):
    row = df.loc[n, ['fn', 'bname', 'label']].values
    audio_path, audio_name, target = row

    if os.path.isfile(audio_path):
      (audio, _) = librosa.core.load(audio_path, sr=sr, mono=True)
      audio = pad_or_truncate(audio, clip_samples)

      w[n] = audio
    else:
      print('{} File does not exist! {}'.format(n, audio_path))

  # Arguments & parameters
  clip_samples = sr*secs
  audios_num = len(df)

  # Pack waveform to hdf5
  total_time = time.time()

  wavs = np.empty((audios_num, clip_samples), dtype=np.float32)
  _ =  Parallel()( delayed(__parallel)(df, wavs, n) for n in tqdm(range(audios_num)) )

  np.save(npy_path, wavs)

  print('Write to {}'.format(npy_path))
  print('Pack npy time: {:.3f}'.format(time.time() - total_time))

  return wavs


# + [markdown] id="qO0pqVlPYVJF"
# ## Blocks functions

# + id="im_DfIXNMfB7"
def init_layer(layer):
  """Initialize a Linear or Convolutional layer. """
  nn.init.xavier_uniform_(layer.weight)

  if hasattr(layer, 'bias'):
    if layer.bias is not None:
      layer.bias.data.fill_(0.)
            
def init_bn(bn):
  """Initialize a Batchnorm layer. """
  bn.bias.data.fill_(0.)
  bn.weight.data.fill_(1.)

class AttBlock(nn.Module):
  def __init__(self, in_features: int, out_features: int, activation="linear", temperature=1.0):
    super().__init__()

    self.activation = activation
    self.temperature = temperature
    self.att = nn.Conv1d(
        in_channels=in_features,
        out_channels=out_features,
        kernel_size=1,
        stride=1,
        padding=0,
        bias=True)
    self.cla = nn.Conv1d(
        in_channels=in_features,
        out_channels=out_features,
        kernel_size=1,
        stride=1,
        padding=0,
        bias=True)

    self.bn_att = nn.BatchNorm1d(out_features)
    self.init_weights()

  def init_weights(self):
    init_layer(self.att)
    init_layer(self.cla)
    init_bn(self.bn_att)

  def forward(self, x):
    # x: (n_samples, n_in, n_time)
    norm_att = torch.softmax(torch.tanh(self.att(x)), dim=-1)
    cla = self.nonlinear_transform(self.cla(x))
    x = torch.sum(norm_att * cla, dim=2)
    return x, norm_att, cla

  def nonlinear_transform(self, x):
    if self.activation == 'linear':
      return x
    elif self.activation == 'sigmoid':
      return torch.sigmoid(x)


# + id="9S8cgekxaqlj"
def get_model(config):
  return PANNsResnetAtt(**config)


# + [markdown] id="ajsVbGzwzhwk"
# ## ResNet

# + id="nWpeRz0xzhwk"
class PANNsResnetAtt(nn.Module):
  def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin, fmax, 
               classes_num, arch='resnet34', fc=512, apply_aug=True, top_db=None, **args):
    super(PANNsResnetAtt, self).__init__()
    
    window = 'hann'
    center = True
    pad_mode = 'reflect'
    ref = 1.0
    amin = 1e-10

    self.interpolate_ratio = 32  # Downsampled ratio
    self.apply_aug = apply_aug

    # Spectrogram extractor
    self.spectrogram_extractor = Spectrogram(
        n_fft=window_size,
        hop_length=hop_size,
        win_length=window_size,
        window=window,
        center=center,
        pad_mode=pad_mode,
        freeze_parameters=True)

    # Logmel feature extractor
    self.logmel_extractor = LogmelFilterBank(
        sr=sample_rate,
        n_fft=window_size,
        n_mels=mel_bins,
        fmin=fmin,
        fmax=fmax,
        ref=ref,
        amin=amin,
        top_db=top_db,
        freeze_parameters=True)

    # Spec augmenter
    self.spec_augmenter = SpecAugmentation(
        time_drop_width=64,
        time_stripes_num=2,
        freq_drop_width=8,
        freq_stripes_num=2)

    self.bn0 = nn.BatchNorm2d(mel_bins)

    att_size = 1024

    self.fc1 = nn.Linear(fc, att_size, bias=True)
    self.att_block = AttBlock(att_size, classes_num, activation='linear')


    resnet = getattr(models, arch)(pretrained=True, progress=False)
    self.resnet_features = nn.Sequential(
        resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool,
        resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
    )
    # del self.resnet_features.avgpool
    # del self.resnet_features.fc

    self.init_weight()

  def init_weight(self):
    init_bn(self.bn0)
    init_layer(self.fc1)
      
  def cnn_feature_extractor(self, x):
    x = self.resnet_features(x)
    return x
  
  def preprocess(self, input, mixup_lambda=None):
    x = self.spectrogram_extractor(input)  # (batch_size, 1, time_steps, freq_bins)
    x = self.logmel_extractor(x)  # (batch_size, 1, time_steps, mel_bins)

    frames_num = x.shape[2]

    x = x.transpose(1, 3)
    x = self.bn0(x)
    x = x.transpose(1, 3)

    if self.training and self.apply_aug:
        x = self.spec_augmenter(x)

    # Mixup on spectrogram
    if self.training  and self.apply_aug and mixup_lambda is not None:
        x = do_mixup(x, mixup_lambda)

    return x, frames_num
      
  def forward(self, input, mixup_lambda=None):
    """
    Input: (batch_size, data_length)"""
    x, frames_num = self.preprocess(input, mixup_lambda=mixup_lambda)

    if mixup_lambda is not None:
        b = (b*c)//2
        c = 1
    
    # Output shape (batch size, channels, time, frequency)
    x = x.expand(x.shape[0], 3, x.shape[2], x.shape[3])
    x = self.cnn_feature_extractor(x)
    
    # Aggregate in frequency axis
    x = torch.mean(x, dim=3)

    x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1)
    x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1)
    x = x1 + x2

    x = F.dropout(x, p=0.5, training=self.training)
    x = x.transpose(1, 2)
    x = F.relu_(self.fc1(x))
    x = x.transpose(1, 2)
    x = F.dropout(x, p=0.5, training=self.training)

    (clipwise_output, norm_att, segmentwise_output) = self.att_block(x)
    # segmentwise_output = segmentwise_output.transpose(1, 2)

    # # Get framewise output
    # framewise_output = interpolate(segmentwise_output,
    #                                 self.interpolate_ratio)
    # framewise_output = pad_framewise_output(framewise_output, frames_num)
    # frame_shape =  framewise_output.shape
    # clip_shape = clipwise_output.shape
    # output_dict = {
    #     'framewise_output': framewise_output.reshape(b, c, frame_shape[1],frame_shape[2]),
    #     'clipwise_output': clipwise_output.reshape(b, c, clip_shape[1]),
    # }

    return clipwise_output


# + [markdown] id="g8yKy5eiZEQd"
# ## Dataset

# + id="WJJ0lMXYnHxt"
class AudioDataset(torch.utils.data.Dataset):
  def __init__(self, df, task='train', **kwargs):
    super(AudioDataset, self).__init__()
    self.df = df
    self.task = task
    self.c = len(words)
    self.classes = words

  def __len__(self):
    return len(self.df)

  def __getitem__(self, idx):
    row = self.df.loc[idx]
    iidx, label = row["index"], 0

    if self.task=='train':
      label = row.label
      waveform = waveforms[iidx]
    else:
      waveform = waveforms_[iidx]

    return {
        'wav': torch.tensor( waveform, dtype=torch.float ),
        'target': torch.tensor( label, dtype=torch.long )
    }


# + [markdown] id="UcPE6p85ZGej"
# ## Training functions

# + id="cnDzEktyS6mH"
def training_fn(dataloader, model, opt, criterion, scheduler=None):
  avg_loss = 0
  avg_acc = 0
  size = len(dataloader)

  model.train()

  for i, data in enumerate(dataloader):
    x,y = data['wav'].to(device), data['target'].to(device)

    opt.zero_grad()

    pred = model(x)
    loss = criterion(pred, y)

    avg_loss += loss.item()
    
    pred = pred.detach().cpu()
    ys = y.detach().cpu()

    avg_acc += (ys == pred.argmax(1)).float().mean().item()

    loss.backward()
    opt.step()
    
    if scheduler:
      scheduler.step()

    print('\r[Training][{}/{}] Loss: {:.5f} - Acc : {:.5f}'.format(
        i+1, size, avg_loss/(i+1), avg_acc/(i+1) ), end='')
  print()
  


# + id="xXwjzFz2S6cL"
def evaluate(dataloader, model, criterion):
  avg_loss = 0
  avg_acc = 0
  size = len(dataloader)

  model.eval()

  with torch.no_grad():
    for i, data in enumerate(dataloader):
      x,y = data['wav'].to(device), data['target'].to(device)

      pred = model(x)
      
      avg_loss += criterion(pred, y).item()

      pred = pred.detach().cpu()
      ys = y.detach().cpu()

      avg_acc += (ys == pred.argmax(1)).float().mean().item()

      print('\r[Evaluation][{}/{}] Loss: {:.5f} - Acc : {:.5f}'.format(
          i+1, size, avg_loss/(i+1), avg_acc/(i+1) ), end='')
    print()
    avg_loss /= size
    avg_acc /= size
    
  return avg_loss


# + id="OHUdrPnUS7CJ"
def predict(df, bs=2):
  test_ds = AudioDataset(df, task='test')
  testloader = torch.utils.data.DataLoader(test_ds, bs, shuffle=False)

  predictions_labels = []
  predictions_proba = []

  out = None

  with torch.no_grad():
    for data in tqdm(testloader):
      x = data['wav'].to(device)

      for i in range(n_folds):
        if i == 0: out = F.softmax( MODELS[i](x), 1 )
        else: out += F.softmax( MODELS[i](x), 1 )

      out /= n_folds
      out_labels = out.argmax(1).cpu().detach().numpy()
      out_probas = out.cpu().detach().numpy()

      predictions_labels += out_labels.tolist()
      predictions_proba += out_probas.tolist()

  return predictions_labels ,predictions_proba


# + id="8QGtjqaxTShX"
def run_fold(fold, config, bs=16, eval_bs=8, lr=1e-4, path='MODELS/'):
  with torch.cuda.device(device):
      torch.cuda.empty_cache()

  best_logloss = np.inf

  fold_train = train[train.fold != fold].reset_index(drop=False)
  fold_val = train[train.fold == fold].reset_index(drop=False)

  train_ds = AudioDataset(fold_train)
  val_ds = AudioDataset(fold_val)

  trainloader = torch.utils.data.DataLoader(train_ds, batch_size=bs, shuffle=True)
  validloader = torch.utils.data.DataLoader(val_ds, batch_size=eval_bs, shuffle=False)

  model = get_model(config)
  criterion = torch.nn.CrossEntropyLoss()
  opt = torch.optim.AdamW(model.parameters(), lr=lr)

  scheduler = None
  if config["schedule"]:
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        opt, max_lr=1e-3, div_factor=4, steps_per_epoch=len(trainloader), epochs=epochs
    )
  
  model.to(device)

  loader = tqdm(range(epochs), desc=f'Fold {fold}')

  for epoch in loader:
    print(f"[Epoch {epoch}]")

    training_fn(trainloader, model, opt, criterion, scheduler)
    avg_logloss = evaluate(validloader, model, criterion)

    if avg_logloss < best_logloss:
      best_logloss = avg_logloss
      torch.save(model.state_dict(), f'{path}model_state_dict_{fold}.bin')

  return best_logloss


# + [markdown] id="x9ZdCZu1Xu3f"
# #Loading the CSVs' files

# + colab={"base_uri": "https://localhost:8080/", "height": 204} id="eSYcXfG6MP6Y" outputId="5bfdadf7-9296-447a-fffb-4809f0712d1a"
train = pd.read_csv(path+'here/Train.csv')
train.head()

# + colab={"base_uri": "https://localhost:8080/", "height": 163} id="I5q58g6xb8Rj" outputId="efab30d5-1efc-4cef-b708-fd1ac8fb6c01"
sub = pd.read_csv(path+'here/SampleSubmission.csv')
sub.head(1)

# + id="Xhp3f-pqMWqb"
words = sub.columns[2:]

# + id="DHJdYHlS_DYY"
label = np.linspace(0, len(words)-1, len(words), dtype=np.int16)
mapper = dict(zip(words, label))

# + id="hnwMYH_O_pvk"
train['label'] = train['target'].map(mapper).astype(int)

# + colab={"base_uri": "https://localhost:8080/", "height": 204} id="A2kk8pEyWEwj" outputId="ee50cd74-b798-43e7-d45f-3f8655800bf7"
train.head()

# + [markdown] id="xi_5mepkmJpE"
# # Save wavs as npy

# + id="BKkVVi7-Crjc"
num_cores = multiprocessing.cpu_count()

# + id="n4A0WzJyHfu9"
sr = 44100
sec = 3

# + id="kwLXyNWePqnw"
npy_path = f'drive/My Drive/Zindi/GIZ/train_sr={sr}_sec={sec}.npy'
test_npy_path = f'drive/My Drive/Zindi/GIZ/test_sr={sr}_sec={sec}.npy'

# + colab={"background_save": true} id="4YzjfYdg9_En"
if not os.path.exists(npy_path):
  waveforms = pack_waveforms_to_npy(npy_path, train, sr=sr, secs=sec)

# + id="bVP1qQOQPS2B"
if not os.path.exists(test_npy_path):
  test = sub[['fn', 'bname']]
  test['label'] = 0

  waveforms_ = pack_waveforms_to_npy(test_npy_path, test, sr=sr, secs=sec)

# + [markdown] id="gioR_yPgtGof"
# # Load wavs

# + colab={"base_uri": "https://localhost:8080/"} id="k7oYBXVnm1g9" outputId="83660894-e0c0-49db-a24c-26acbc17a7fc"
# %%time
waveforms = load_npy(npy_path)
waveforms_ = load_npy(test_npy_path)

# + colab={"base_uri": "https://localhost:8080/"} id="n5O6I6LCYcdz" outputId="87e64d8e-fb16-4f9b-ed25-445ddaeb6ea1"
gc.collect()

# + [markdown] id="IKQ5aLUMmQbV"
# #Training

# + id="Uggl3ojlVeEz"
n_folds = 10


# + id="7Mh4eTUZI86H"
def make_fold(n_folds, shuffle=True):

  train['fold'] = -1
  indexes = train[train['type']=='base'].index
  
  fold = StratifiedKFold(n_splits = n_folds, random_state=seed, shuffle=shuffle)
  for i, (tr, vr) in enumerate(fold.split(indexes, train.loc[indexes, 'label'])):
    train.loc[indexes[vr], 'fold'] = i


# + id="xdtoz-WvS_Y0"
make_fold(n_folds, shuffle=True)

# + colab={"base_uri": "https://localhost:8080/"} id="SVBpTXVGLfqA" outputId="f22e8847-5476-4ff8-9f56-cb5b012f0c33"
train.fold.nunique()

# + id="0b1VNOENO4UW"
epochs = 50
device = 'cuda:0'
lr = 1e-4

classes_num = 193
batch_size = 16

config = {
    "sample_rate": sr,
    "window_size": 1024,
    "hop_size": 320,
    "mel_bins": 64,
    "fmin": 50,
    "fmax": 14000,
    "classes_num": classes_num,
    'arch': 'resnet50',
    'fc': 2048,
    "schedule": False,
}

# + colab={"base_uri": "https://localhost:8080/"} id="w8fUyfUtKSc1" outputId="4c3a4dfc-b50a-4695-a475-f6504addc914"
gc.collect()

# + id="kT0fGu2j1Oiz"
# %%time
avg_logloss = 0

for fold in range(n_folds):
  
  _fold_logloss = run_fold(fold, config, bs=batch_size, eval_bs=batch_size, lr=lr)
  avg_logloss += _fold_logloss

print()
print("Avg LogLoss: ", avg_logloss/n_folds)
print()

# + [markdown] id="BTT9_PEmRTJr"
# #Loading models

# + colab={"base_uri": "https://localhost:8080/"} id="sbdLwC2f2ISf" outputId="af2b80a6-b082-4574-8848-df6b07b00563"
# %%time
MODELS = []

for i in range(n_folds):
  MODELS.append( get_model(config) )
  MODELS[i].load_state_dict(torch.load(f'MODELS/model_state_dict_{i}.bin'))
  MODELS[i].to(device)
  MODELS[i].eval()

# + [markdown] id="9LDMrBFeZcFt"
# #Prediction

# + colab={"base_uri": "https://localhost:8080/", "height": 66, "referenced_widgets": ["414a8a0814664d25aa00a88a23dc4e86"]} id="3ctfz162b5fx" outputId="8cfcb263-0a6d-4248-a030-96bd53263908"
predictions_labels, predictions_proba = predict(sub.reset_index(), bs=2)

# + [markdown] id="yj9LGWbBq1kg"
# # Making a submission

# + id="CzPa_Qqae3LC"
submission = pd.DataFrame()
submission['fn'] = sub['fn'].apply(lambda x: '/'.join( x.split('/')[1:] ))
for i, label in enumerate(words):
  submission[label] = 0.
for (label, i) in mapper.items():
  submission.loc[:,label] = np.array(predictions_proba)[:,i]

# + id="tvGMYQKihpi1"
submission.head()

# + id="2cFo2HSCiA_w"
csv_file = 'resnet50_base.csv'
submission.to_csv(csv_file, index=False)

StarcoderdataJupyterScriptsDedupFiltered