script.py

import io
import time
import os

os.environ['OMP_NUM_THREADS'] = '1'
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
os.environ['NUMEXPR_NUM_THREADS'] = '1'

import torch
import tqdm.auto as tqdm
from glob import glob
import numpy as np
import pandas as pd

from models import Model
from datasets import load_dataset
from preprocess import preprocess, preprocess_FS

from src.rawnet_model import RawNet
from src.lcnn_model import LCNN
from src.resnet_model import ResNet_LogSpec, ResNet_MelSpec
from src.moe_model import UltimateMOE, MOE_attention, MOE_attention_FS

# Import your model and anything else you want
# You can even install other packages included in your repo
# However, during the evaluation the container will not have access to the internet.
# So you must include everything you need in your model repo. Common python libraries will be installed.
# Feel free to contact us to add dependencies to the requiremnts.txt
# For testing, this is the docker image that will be used https://github.com/huggingface/competitions/blob/main/Dockerfile
# It can be pulled here https://hub.docker.com/r/huggingface/competitions/tags


# load the dataset. dataset will be automatically downloaded to /tmp/data during evaluation
print('Load Dataset')
# DATASET_PATH = "/tmp/data_test"
# dataset_remote = glob(os.path.join(DATASET_PATH, '*'))
DATASET_PATH = "/tmp/data"
dataset_remote = load_dataset(DATASET_PATH,split="test", streaming=True)

device = "cuda:0"
# device = "cpu"

print('Define Model')

# # RAWNET2 MODEL
# config = {
#     "first_conv": 1024, "in_channels": 1, "filts": [20, [20, 20], [20, 128], [128, 128]],
#     "blocks": [2, 4], "nb_fc_node": 1024, "gru_node": 1024, "nb_gru_layer": 3, "nb_classes": 2
# }
# model = RawNet(config, device).to(device)
# model_path = './checkpoints/RAWNET_ASVSPOOF_FOR_INTHEWILD_PURDUE.pth'
# model.load_state_dict(torch.load(model_path, map_location=device))

# RESNET MODEL
# model = ResNet_LogSpec(sample_rate=24000, return_emb=False).to(device)
# model_path = './checkpoints/RESNET_LOGSPEC_ALL_DATA_FS_24000.pth'

# model = ResNet_MelSpec(sample_rate=24000, return_emb=False).to(device)
# model_path = './checkpoints/RESNET_MELSPEC_ALL_DATA_FS_24000.pth'

## LCNN MODEL
# model = LCNN(return_emb=False, fs=24000).to(device)
# model_path = './checkpoints/LCNN_ASVSPOOF_FOR_INTHEWILD_PURDUE.pth'
# model_path = './checkpoints/LCNN_ALL_DATA.pth'
# model_path = './checkpoints/LCNN_ALL_DATA_AUG.pth'
# model_path = './checkpoints/LCNN_ALL_DATA_TTS_AUG.pth'
# model_path = './checkpoints/LCNN_ALL_DATA_TTS_MOD.pth'
# model_path = './checkpoints/LCNN_ALL_DATA_HI_FREQ_22050.pth'

# model_path = './checkpoints/LCNN_ALL_DATA_FS_16000.pth'
# model_path = './checkpoints/LCNN_ALL_DATA_FS_22050.pth'
# model_path = './checkpoints/LCNN_ALL_DATA_FS_24000.pth'

# model.load_state_dict(torch.load(model_path, map_location=device))

# # MOE MODEL
expert_1 = LCNN(return_emb=True, fs=24000)
expert_2 = ResNet_LogSpec(return_emb=True, sample_rate=24000)
expert_3 = ResNet_MelSpec(return_emb=True, sample_rate=24000)

model = MOE_attention(experts=[expert_1, expert_2, expert_3], device=device)
model_path = './checkpoints/MOE_TRANSF_3EXP_MODELS_AUG.pth'

# expert_1 = LCNN(return_emb=True, fs=16000).to(device)
# expert_2 = LCNN(return_emb=True, fs=22050).to(device)
# expert_3 = LCNN(return_emb=True, fs=24000).to(device)
# # expert_4 = LCNN(return_emb=True).to(device)
# # expert_5 = LCNN(return_emb=True).to(device)
# # expert_6 = LCNN(return_emb=True).to(device)
#
# model = MOE_attention_FS(experts=[expert_1, expert_2, expert_3], device=device)
# model_path = './checkpoints/MOE_TRANSF_3EXP_FS_AUG_NO_FREEZE.pth'

# # # model = UltimateMOE(experts=[expert_1, expert_2, expert_3, expert_4])
# # # model_path = './checkpoints/MOE_ULTIMATE.pth'
#
# # model = MOE_attention(experts=[expert_1, expert_2, expert_3, expert_4, expert_5, expert_6], device=device)
# # # model_path = './checkpoints/MOE_ATTENTION.pth'
# # model_path = './checkpoints/MOE_TRANSF.pth'
#
# expert_7 = LCNN(return_emb=True).to(device)
# expert_8 = LCNN(return_emb=True).to(device)
# model = MOE_attention(experts=[expert_1, expert_2, expert_3, expert_4, expert_5, expert_6, expert_7, expert_8], device=device, freezing=True)
# # model_path = './checkpoints/MOE_TRANSF_7EXP.pth'
# # model_path = './checkpoints/MOE_TRANSF_7EXP_AUG.pth'
# # model_path = './checkpoints/MOE_TRANSF_7EXP_AUG_NO_FREEZE.pth'
# # model_path = './checkpoints/MOE_TRANSF_8EXP_AUG.pth'
# model_path = './checkpoints/MOE_TRANSF_8EXP_AUG_NO_FREEZE.pth'

model = (model).to(device)
model.load_state_dict(torch.load(model_path, map_location=device))

model.eval()

print('Loaded Weights')

# # EVALUATE OLD MODEL
# del model
# model = Model().to(device)

# SAMPLING_RATE_CODES = {
#     8000: 2,
#     16000: 3,
#     22050: 5,
#     24000: 7,
#     32000: 11,
#     44100: 13,
#     48000: 17,
#     "other": 19
# }
#
# seen_frequencies = set()

# iterate over the dataset
out = []
for el in tqdm.tqdm(dataset_remote):
# for el in dataset_remote:

    start_time = time.time()

    # each element is a dict
    # el["id"] id of example and el["audio"] contains the audio file
    # el["audio"]["bytes"] contains bytes from reading the raw audio
    # el["audio"]["path"] containts the filename. This is just for reference and you cant actually load it

    # if you are using libraries that expect a file. You can use BytesIO object

    # try:

    # RUNNING ON HUGGINGFACE
    file_like = io.BytesIO(el["audio"]["bytes"])

    tensor, sr = preprocess(file_like, target_sr=24000)
    # tensor_16, tensor_22, tensor_24 = preprocess_FS(file_like)

    # # RUNNING LOCALLY
    # tensor = preprocess(el)

    with torch.no_grad():
        # soft decision (such as log likelihood score)
        # positive score correspond to synthetic prediction
        # negative score correspond to pristine prediction

        # # OLD MODEL
        # score = model(tensor.to(device)).cpu().item()

        # CUSTOM MODEL
        score = model(tensor.to(device))[:, 1].cpu()
        # score = model(tensor_16.to(device), tensor_22.to(device), tensor_24.to(device))[:, 1].cpu()

        print(f'SCORE OUT: {score}')
        score = score.mean().item()
        print(f'SCORE FINAL: {score}')

        # we require a hard decision to be submited. so you need to pick a threshold
        pred = "generated" if score > model.threshold else "pristine"

    # append your prediction
    # "id" and "pred" are required. "score" will not be used in scoring but we encourage you to include it. We'll use it for analysis of the results

    # RUNNING ON HUGGINGFACE
    total_time = time.time() - start_time

    # freq = sr if sr in SAMPLING_RATE_CODES else "other"
    #
    # # Assegna total_time: codice se è la prima occorrenza, 0 altrimenti
    # if freq not in seen_frequencies:
    #     total_time = SAMPLING_RATE_CODES[freq]
    #     seen_frequencies.add(freq)
    # # else:
    # #     total_time = 0
    # total_time = 1

    out.append(dict(id=el["id"], pred=pred, score=score, time=total_time))
    # # RUNNING LOCALLY
    # out.append(dict(id=el, pred=pred, score=score, time=time.time() - start_time))

    # except Exception as e:
    #     print(e)
    #     print("failed", el["id"])
    #     out.append(dict(id=el["id"], pred="none", score=None))
    #     # print("failed", el)
    #     # out.append(dict(id=el, pred="none", score=None))

# save the final result and that's it
pd.DataFrame(out).to_csv("submission.csv", index=False)