Repository: LAION-AI/CLAP (language: python)

Issue:
pip merge with clap branch



Discussion / hints:

Fix patch:
diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 0000000..ee6efa5
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1,3 @@
+recursive-include src/laion_clap/clap_module/model_configs *.json
+recursive-include src/laion_clap/clap_module bpe_simple_vocab_16e6.txt.gz
+recursive-include src/laion_clap/training audioset_textmap.npy
diff --git a/README.md b/README.md
index 8745d5c..f5499c1 100644
--- a/README.md
+++ b/README.md
@@ -2,7 +2,7 @@
 
 Contrastive Language-Audio Pretraining, known as CLAP. Referring to the CLIP (Contrastive Language-Image Pretraining) architecture, similarly, the CLAP architecture is as follows.  
 <p align="center">
-  <img src="./assets/audioclip-arch.png" alt="The Contrastive Language-Audio Pretraining Model Architecture" width="60%"/>
+  <img src="https://raw.githubusercontent.com/LAION-AI/CLAP/main/assets/audioclip-arch.png" alt="The Contrastive Language-Audio Pretraining Model Architecture" width="60%"/>
 </p>
 
 
@@ -18,8 +18,47 @@ The major opensource contributers of this project are (in equal contribution): Y
 
 many thanks to <a href="https://github.com/cfoster0/CLAP">@cfoster0</a> for allowing us to use his repo name.
 
+## Quick Start 
+We provide the library for our CLAP model:
+```bash
+pip install laion_clap
+```
+
+Then you can follow the below usage or refer to [unit_test.py](https://github.com/LAION-AI/CLAP/blob/laion_clap_pip/src/laion_clap/unit_test.py).
+
+```python
+import librosa
+import laion_clap
+
+model = laion_clap.CLAP_Module(enable_fusion=True)
+model.load_ckpt()
+
+# Directly get audio embeddings from audio files
+audio_file = [
+    '/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_short.wav',
+    '/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_long.wav'
+]
+audio_embed = model.get_audio_embedding_from_filelist(x = audio_file)
+print(audio_embed)
+print(audio_embed.shape)
+
+# Get audio embeddings from audio data
+audio_data, _ = librosa.load('/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_short.wav', sr=48000) # sample rate should be 48000
+audio_data = audio_data.reshape(1, -1) # Make it (1,T) or (N,T)
+
+audio_embed = model.get_audio_embedding_from_data(x = audio_data)
+print(audio_embed)
+print(audio_embed.shape)
+
+# Get text embedings from texts:
+text_data = ["I love the contrastive learning", "I love the pretrain model"] 
+text_embed = model.get_text_embedding(text_data)
+print(text_embed)
+print(text_embed.shape)
+```
+
 ## Environment Installation
-To install the same environment as we use, please run the following command:
+If you want to check and reuse our model into your project instead of directly using the pip library, you need to install the same environment as we use, please run the following command:
 ```bash
 conda create env -n clap python=3.10
 conda activate clap
@@ -43,52 +82,11 @@ For example, in a single machine multi-GPU setting, you might want to use `torch
 To train on a single GPU machine, use `CUDA_VISIBLE_DEVICES=0 python -m ...` instead of `srun`.
 We use [Weights and Biases](https://wandb.ai/site) for experiment logging. You need to configure the weights and biases in your environment.
 
-## Loading Model and Inference
-Please refer to [infer_demo.py](src/training/infer_demo.py) to get the whole view of using our model to infer the audio and text embeddings.
-Below is the core code.
-```python
-# import necessary libraries
-def infer_audio():
-    
-    '''
-    set hyperparameters, and load pretrain model
-    '''
-    
-    # load the waveform of the shape (T,), should resample to 48000
-    audio_waveform, sr = librosa.load('/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_long.wav', sr=48000) 
-    # quantize
-    audio_waveform = int16_to_float32(float32_to_int16(audio_waveform))
-    audio_waveform = torch.from_numpy(audio_waveform).float()
-    audio_dict = {}
-
-    # the 'fusion' truncate mode can be changed to 'rand_trunc' if run in unfusion mode
-    audio_dict = get_audio_features(
-        audio_dict, audio_waveform, 480000, 
-        data_truncating='fusion', 
-        data_filling='repeatpad',
-        audio_cfg=model_cfg['audio_cfg']
-    )
-    # can send a list to the model, to process many audio tracks in one time (i.e. batch size)
-    audio_embed = model.get_audio_embedding([audio_dict])
-    print(audio_embed.size())
-
-def infer_text():
-    '''
-    set hyperparameters, and load pretrain model
-    '''
-    
-    # load the text, can be a list (i.e. batch size)
-    text_data = ["I love the contrastive learning", "I love the pretrain model"] 
-    # tokenize for roberta, if you want to tokenize for another text encoder, please refer to data.py#L43-90 
-    text_data = tokenizer(text_data)
-    
-    text_embed = model.get_text_embedding(text_data)
-    print(text_embed.size())
-    
-```
+## Core Code
+Please refer to [main.py](https://github.com/LAION-AI/CLAP/blob/laion_clap_pip/src/laion_clap/training/main.py), [train.py](https://github.com/LAION-AI/CLAP/blob/laion_clap_pip/src/laion_clap/training/train.py), [data.py](https://github.com/LAION-AI/CLAP/blob/laion_clap_pip/src/laion_clap/training/data.py),and [model.py](https://github.com/LAION-AI/CLAP/blob/laion_clap_pip/src/laion_clap/clap_module/model.py) to quicly get familiar with our model.
 
 ## Pretrained Models
-The pretrained checkpoints can be found in [here](https://drive.google.com/drive/folders/1Ni8lZ2pryTESjgq8gELLQNM_HGdWtFrE?usp=sharing).
+The pretrained checkpoints can be found in [here](https://huggingface.co/lukewys/laion_clap/tree/main).
 Please refer to the previous section for how to load and run the checkpoints.
 
 The checkpoints list here for each model setting is the one with the highest average mAP score in training.
@@ -106,7 +104,13 @@ If you find this project and the LAION-Audio-630K dataset useful, please cite ou
   title = {Large-scale Contrastive Language-Audio Pretraining with Feature Fusion and Keyword-to-Caption Augmentation},
   author = {Wu*, Yusong and Chen*, Ke and Zhang*, Tianyu and Hui*, Yuchen and Berg-Kirkpatrick, Taylor and Dubnov, Shlomo},
   booktitle={IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP},
-  year = {2023},
+  year = {2023}
+}
+@inproceedings{htsatke2022,
+  author = {Ke Chen and Xingjian Du and Bilei Zhu and Zejun Ma and Taylor Berg-Kirkpatrick and Shlomo Dubnov},
+  title = {HTS-AT: A Hierarchical Token-Semantic Audio Transformer for Sound Classification and Detection},
+  booktitle={IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP},
+  year = {2022}
 }
 ```
 
diff --git a/experiment_scripts/esc50_api.py b/experiment_scripts/esc50_api.py
new file mode 100644
index 0000000..37ac0ad
--- /dev/null
+++ b/experiment_scripts/esc50_api.py
@@ -0,0 +1,48 @@
+import laion_clap
+from tqdm import tqdm
+import glob
+import json
+import torch
+import numpy as np
+
+device = torch.device('cuda:0')  # ('cuda:0')
+
+esc50_test_dir = './ESC50_1/test/*/'
+class_index_dict_path = '/fsx/yusong/CLAP/class_labels/ESC50_class_labels_indices_space.json'
+
+# Load the model
+model = laion_clap.CLAP_Module()
+model.load_ckpt()
+
+# Get the class index dict
+class_index_dict = {v: k for v, k in json.load(open(class_index_dict_path)).items()}
+
+# Get all the data
+audio_files = sorted(glob.glob(esc50_test_dir + '**/*.flac', recursive=True))
+json_files = sorted(glob.glob(esc50_test_dir + '**/*.json', recursive=True))
+ground_truth_idx = [class_index_dict[json.load(open(jf))['tag'][0]] for jf in json_files]
+
+with torch.no_grad():
+    ground_truth = torch.tensor(ground_truth_idx).view(-1, 1)
+
+    # Get text features
+    all_texts = ["This is a sound of " + t for t in class_index_dict.keys()]
+    text_embed = model.get_text_embedding(all_texts)
+    audio_embed = model.get_audio_embedding_from_filelist(x=audio_files)
+
+    ranking = torch.argsort(torch.tensor(audio_embed) @ torch.tensor(text_embed).t(), descending=True)
+    preds = torch.where(ranking == ground_truth)[1]
+    preds = preds.cpu().numpy()
+
+    metrics = {}
+    metrics[f"mean_rank"] = preds.mean() + 1
+    metrics[f"median_rank"] = np.floor(np.median(preds)) + 1
+    for k in [1, 5, 10]:
+        metrics[f"R@{k}"] = np.mean(preds < k)
+    # map@10
+    metrics[f"mAP@10"] = np.mean(np.where(preds < 10, 1 / (preds + 1), 0.0))
+
+    print(
+        f"Zeroshot Classification Results: "
+        + "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in metrics.items()])
+    )
diff --git a/experiment_scripts/zeroshot_esc50.sh b/experiment_scripts/zeroshot_esc50.sh
index 9512548..cfba10f 100644
--- a/experiment_scripts/zeroshot_esc50.sh
+++ b/experiment_scripts/zeroshot_esc50.sh
@@ -1,64 +1,105 @@
-#!/bin/bash
-#SBATCH --comment clap
-#SBATCH --partition=g40423
-#SBATCH --job-name=mclap
-#SBATCH --nodes 3
-#SBATCH --ntasks-per-node 8
-#SBATCH --cpus-per-gpu=6
-#SBATCH --exclusive
-#SBATCH --output=%x_%j.out
+# run from CLAP directory
+python -m evaluate.eval_zeroshot_classification \
+  --dataset-type="webdataset" \
+  --precision="fp32" \
+  --batch-size=512 \
+  --workers=6 \
+  --amodel HTSAT-tiny \
+  --tmodel roberta \
+  --datasetnames "esc50_no_overlap" \
+  --remotedata \
+  --datasetinfos "train" \
+  --seed 3407 \
+  --logs ./logs \
+  --data-filling "repeatpad" \
+  --data-truncating "rand_trunc" \
+  --freeze-text \
+  --class-label-path="../class_labels/ESC50_class_labels_indices_space.json" \
+  --pretrained="/fsx/clap_logs/2023_02_18-00_03_45-model_HTSAT-tiny-lr_0.0001-b_96-j_6-p_fp32/checkpoints"
 
-module load openmpi
-module load cuda/11.7
-export NCCL_PROTO=simple
-export FI_EFA_FORK_SAFE=1
-export FI_LOG_LEVEL=1
-export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
-export NCCL_DEBUG=info
-export OMPI_MCA_mtl_base_verbose=1
-export FI_EFA_ENABLE_SHM_TRANSFER=0
-export FI_PROVIDER=efa
-export FI_EFA_TX_MIN_CREDITS=64
-export NCCL_TREE_THRESHOLD=0
 
-# sent to sub script
-export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
-export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
-export MASTER_PORT=12802
-export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+python -m evaluate.eval_zeroshot_classification \
+  --dataset-type="webdataset" \
+  --precision="fp32" \
+  --batch-size=512 \
+  --workers=6 \
+  --amodel HTSAT-tiny \
+  --tmodel roberta \
+  --report-to "wandb" \
+  --wandb-notes "2.26-zeroshot-esc50-no-overlap-dataset-5#-fusion" \
+  --datasetnames "esc50_no_overlap" \
+  --remotedata \
+  --datasetinfos "train" \
+  --seed 3407 \
+  --logs ./logs \
+  --data-filling "repeatpad" \
+  --data-truncating "fusion" \
+  --enable-fusion \
+  --fusion-type "aff_2d" \
+  --freeze-text \
+  --class-label-path="../class_labels/ESC50_class_labels_indices_space.json" \
+  --pretrained="/fsx/clap_logs/2023_02_26-01_28_30-model_HTSAT-tiny-lr_0.0001-b_96-j_6-p_fp32/checkpoints"
 
-echo go $COUNT_NODE
-echo $HOSTNAMES
 
-source /fsx/yusong/clap/bin/activate
-cd /fsx/yusong/CLAP/src
-export TRANSFORMERS_CACHE=/fsx/yusong/transformers_cache
+python -m evaluate.eval_zeroshot_classification \
+  --dataset-type="webdataset" \
+  --precision="fp32" \
+  --batch-size=512 \
+  --workers=4 \
+  --amodel HTSAT-tiny \
+  --tmodel roberta \
+  --report-to "wandb" \
+  --wandb-notes "2.26-zeroshot-esc50-no-overlap-dataset-2#-fusion" \
+  --datasetnames "esc50_no_overlap" \
+  --remotedata \
+  --datasetinfos "train" \
+  --seed 3407 \
+  --logs ./logs \
+  --data-filling "repeatpad" \
+  --data-truncating "fusion" \
+  --enable-fusion \
+  --fusion-type "aff_2d" \
+  --freeze-text \
+  --class-label-path="../class_labels/ESC50_class_labels_indices_space.json" \
+  --pretrained="/fsx/clap_logs/2023_02_26-01_28_31-model_HTSAT-tiny-lr_0.0001-b_96-j_6-p_fp32/checkpoints"
 
-srun --comment clap --cpu_bind=v --accel-bind=gn python -m evaluate.eval_retrieval_main \
-    --save-frequency 5 \
-    --save-top-performance 3 \
-    --save-most-recent \
-    --dataset-type="webdataset" \
-    --precision="fp32" \
-    --warmup 0 \
-    --batch-size=512 \
-    --wd=0.0 \
-    --epochs=50 \
-    --workers=6 \
-    --use-bn-sync \
-    --freeze-text \
-    --amodel HTSAT-tiny \
-    --tmodel roberta \
-    --report-to "wandb" \
-    --wandb-notes "10.17-zeroshot-esc50-dataset-4#" \
-    --datasetnames "esc50" \
-    --datasetinfos "train" \
-    --seed 3407 \
-    --remotedata \
-    --logs /fsx/clap_logs \
-    --gather-with-grad \
-    --openai-model-cache-dir /fsx/yusong/transformers_cache \
-    --data-filling "repeatpad" \
-    --data-truncating "rand_trunc" \
-    --class-label-path="../class_labels/ESC50_class_labels_indices_space.json" \
-    --pretrained="/fsx/clap_logs/2022_10_17-02_08_21-model_HTSAT-tiny-lr_0.0001-b_96-j_6-p_fp32/checkpoints"
\ No newline at end of file
+
+# 90.3% accuracy
+python -m evaluate.eval_zeroshot_classification \
+  --dataset-type="webdataset" \
+  --precision="fp32" \
+  --batch-size=512 \
+  --workers=6 \
+  --amodel HTSAT-tiny \
+  --tmodel roberta \
+  --datasetnames "ESC50_1" \
+  --remotedata \
+  --datasetinfos "train" \
+  --seed 3407 \
+  --logs ./logs \
+  --data-filling "repeatpad" \
+  --data-truncating "rand_trunc" \
+  --freeze-text \
+  --class-label-path="../class_labels/ESC50_class_labels_indices_space.json" \
+  --pretrained="/fsx/clap_logs/2022_10_17-02_08_30-model_HTSAT-tiny-lr_0.0001-b_96-j_6-p_fp32/checkpoints/epoch_15.pt"
+
+
+python -m evaluate.eval_zeroshot_classification \
+--dataset-type="webdataset" \
+--precision="fp32" \
+--batch-size=512 \
+--workers=4 \
+--amodel HTSAT-tiny \
+--tmodel roberta \
+--datasetnames "ESC50_1" \
+--remotedata \
+--datasetinfos "train" \
+--seed 3407 \
+--logs ./logs \
+--data-filling "repeatpad" \
+--data-truncating "fusion" \
+--enable-fusion \
+--fusion-type "aff_2d" \
+--freeze-text \
+--class-label-path="../class_labels/ESC50_class_labels_indices_space.json" \
+--pretrained="/fsx/clap_logs/2023_02_26-01_28_30-model_HTSAT-tiny-lr_0.0001-b_96-j_6-p_fp32/checkpoints/epoch_20.pt"
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..1be095a
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,53 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+[project]
+name = "laion_clap"
+version = "1.0.0"
+authors = [
+  { name="Ke Chen", email="knutchen@ucsd.edu" },
+  { name="Yusong Wu" },
+  { name="Tianyu Zhang" },
+  { name="Yuchen Hui" }
+]
+maintainers = [
+  { name="Ke Chen", email="knutchen@ucsd.edu" },
+  { name="Yusong Wu" },
+  { name="Tianyu Zhang" },
+  { name="Yuchen Hui" }
+]
+description = "Contrastive Language-Audio Pretraining Model from LAION"
+license = {file = "LICENSE"}
+readme = "README.md"
+requires-python = ">=3.7"
+dependencies = [
+  "soundfile",
+  "librosa",
+  "torchlibrosa",
+  "ftfy",
+  "braceexpand",
+  "webdataset",
+  "wget",
+  "wandb",
+  "llvmlite",
+  "scipy",
+  "scikit-learn",
+  "pandas",
+  "h5py",
+  "tqdm",
+  "regex",
+  "transformers",
+  "progressbar"
+]
+classifiers = [
+    'Development Status :: 3 - Alpha',
+    'Intended Audience :: Developers',
+    'Intended Audience :: Science/Research',
+    'License :: OSI Approved :: Apache Software License',
+    'Topic :: Scientific/Engineering :: Artificial Intelligence',
+]
+
+
+[project.urls]
+"Homepage" = "https://github.com/LAION-AI/CLAP"
+"Bug Tracker" = "https://github.com/LAION-AI/CLAP/issues"
\ No newline at end of file
diff --git a/src/bootstrap_pytorch_dist_env.sh b/src/bootstrap_pytorch_dist_env.sh
deleted file mode 100755
index df6c652..0000000
--- a/src/bootstrap_pytorch_dist_env.sh
+++ /dev/null
@@ -1,17 +0,0 @@
-#!/bin/bash
-# Set env vars for PyTorch
-nodes=($(cat ${LSB_DJOB_HOSTFILE} | sort | uniq | grep -v login | grep -v batch))
-head=${nodes[0]}
-
-export RANK=$OMPI_COMM_WORLD_RANK
-export LOCAL_RANK=$OMPI_COMM_WORLD_LOCAL_RANK
-export WORLD_SIZE=$OMPI_COMM_WORLD_SIZE
-export MASTER_ADDR=$head
-export MASTER_PORT=29500 # default from torch launcher
-
-echo "Setting env_var RANK=${RANK}"
-echo "Setting env_var LOCAL_RANK=${LOCAL_RANK}"
-echo "Setting env_var WORLD_SIZE=${WORLD_SIZE}"
-echo "Setting env_var MASTER_ADDR=${MASTER_ADDR}"
-echo "Setting env_var MASTER_PORT=${MASTER_PORT}"
-
diff --git a/src/laion_clap/__init__.py b/src/laion_clap/__init__.py
new file mode 100644
index 0000000..43420fe
--- /dev/null
+++ b/src/laion_clap/__init__.py
@@ -0,0 +1,5 @@
+import os
+import sys
+dir_path = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(dir_path)
+from hook import CLAP_Module
\ No newline at end of file
diff --git a/src/open_clip/__init__.py b/src/laion_clap/clap_module/__init__.py
similarity index 93%
rename from src/open_clip/__init__.py
rename to src/laion_clap/clap_module/__init__.py
index 96ccf3e..b585be6 100644
--- a/src/open_clip/__init__.py
+++ b/src/laion_clap/clap_module/__init__.py
@@ -5,4 +5,4 @@
 from .pretrained import list_pretrained, list_pretrained_tag_models, list_pretrained_model_tags,\
     get_pretrained_url, download_pretrained
 from .tokenizer import SimpleTokenizer, tokenize
-from .transform import image_transform
+from .transform import image_transform
\ No newline at end of file
diff --git a/src/open_clip/bert.py b/src/laion_clap/clap_module/bert.py
similarity index 100%
rename from src/open_clip/bert.py
rename to src/laion_clap/clap_module/bert.py
diff --git a/src/open_clip/bpe_simple_vocab_16e6.txt.gz b/src/laion_clap/clap_module/bpe_simple_vocab_16e6.txt.gz
similarity index 100%
rename from src/open_clip/bpe_simple_vocab_16e6.txt.gz
rename to src/laion_clap/clap_module/bpe_simple_vocab_16e6.txt.gz
diff --git a/src/open_clip/factory.py b/src/laion_clap/clap_module/factory.py
similarity index 100%
rename from src/open_clip/factory.py
rename to src/laion_clap/clap_module/factory.py
diff --git a/src/open_clip/feature_fusion.py b/src/laion_clap/clap_module/feature_fusion.py
similarity index 100%
rename from src/open_clip/feature_fusion.py
rename to src/laion_clap/clap_module/feature_fusion.py
diff --git a/src/open_clip/htsat.py b/src/laion_clap/clap_module/htsat.py
similarity index 98%
rename from src/open_clip/htsat.py
rename to src/laion_clap/clap_module/htsat.py
index db96116..bb8e7cf 100644
--- a/src/open_clip/htsat.py
+++ b/src/laion_clap/clap_module/htsat.py
@@ -867,8 +867,17 @@ def forward(self, x: torch.Tensor, mixup_lambda = None, infer_mode = False, devi
 
         if self.enable_fusion and x["longer"].sum() == 0:
             # if no audio is longer than 10s, then randomly select one audio to be longer
-            x["longer"][torch.randint(0, x["longer"].shape[0], (1,))] = True
-
+            if self.training:
+                x["longer"][torch.randint(0, x["longer"].shape[0], (1,))] = True
+            else:
+                x = x["mel_fusion"].to(device=device, non_blocking=True)
+                x = x.transpose(1, 3)
+                x = self.bn0(x)
+                x = x.transpose(1, 3)
+                x = self.reshape_wav2img(x)
+                output_dict = self.forward_features(x, longer_idx=[])
+                return output_dict
+                
         if not self.enable_fusion:
             x = x["waveform"].to(device=device, non_blocking=True)
             x = self.spectrogram_extractor(x)   # (batch_size, 1, time_steps, freq_bins)
diff --git a/src/open_clip/linear_probe.py b/src/laion_clap/clap_module/linear_probe.py
similarity index 100%
rename from src/open_clip/linear_probe.py
rename to src/laion_clap/clap_module/linear_probe.py
diff --git a/src/open_clip/loss.py b/src/laion_clap/clap_module/loss.py
similarity index 100%
rename from src/open_clip/loss.py
rename to src/laion_clap/clap_module/loss.py
diff --git a/src/open_clip/model.py b/src/laion_clap/clap_module/model.py
similarity index 99%
rename from src/open_clip/model.py
rename to src/laion_clap/clap_module/model.py
index e35a2c7..60663ec 100644
--- a/src/open_clip/model.py
+++ b/src/laion_clap/clap_module/model.py
@@ -736,10 +736,9 @@ def get_audio_embedding(self, data):
         keys = data[0].keys()
         for k in keys:
             input_dict[k] = torch.cat([d[k].unsqueeze(0) for d in data], dim=0).to(device)
-        
-        audio_embeds = self.audio_projection(self.encode_audio(input_dict, device=device)["embedding"])
+        audio_embeds = self.encode_audio(input_dict, device=device)["embedding"]
+        audio_embeds = self.audio_projection(audio_embeds)
         audio_embeds = F.normalize(audio_embeds, dim=-1)
-        
         return audio_embeds
 
             
diff --git a/src/open_clip/model_configs/HTSAT-base.json b/src/laion_clap/clap_module/model_configs/HTSAT-base.json
similarity index 100%
rename from src/open_clip/model_configs/HTSAT-base.json
rename to src/laion_clap/clap_module/model_configs/HTSAT-base.json
diff --git a/src/open_clip/model_configs/HTSAT-large.json b/src/laion_clap/clap_module/model_configs/HTSAT-large.json
similarity index 100%
rename from src/open_clip/model_configs/HTSAT-large.json
rename to src/laion_clap/clap_module/model_configs/HTSAT-large.json
diff --git a/src/open_clip/model_configs/HTSAT-tiny-win-1536.json b/src/laion_clap/clap_module/model_configs/HTSAT-tiny-win-1536.json
similarity index 100%
rename from src/open_clip/model_configs/HTSAT-tiny-win-1536.json
rename to src/laion_clap/clap_module/model_configs/HTSAT-tiny-win-1536.json
diff --git a/src/open_clip/model_configs/HTSAT-tiny.json b/src/laion_clap/clap_module/model_configs/HTSAT-tiny.json
similarity index 100%
rename from src/open_clip/model_configs/HTSAT-tiny.json
rename to src/laion_clap/clap_module/model_configs/HTSAT-tiny.json
diff --git a/src/open_clip/model_configs/PANN-10.json b/src/laion_clap/clap_module/model_configs/PANN-10.json
similarity index 100%
rename from src/open_clip/model_configs/PANN-10.json
rename to src/laion_clap/clap_module/model_configs/PANN-10.json
diff --git a/src/open_clip/model_configs/PANN-14-fmax-18k.json b/src/laion_clap/clap_module/model_configs/PANN-14-fmax-18k.json
similarity index 100%
rename from src/open_clip/model_configs/PANN-14-fmax-18k.json
rename to src/laion_clap/clap_module/model_configs/PANN-14-fmax-18k.json
diff --git a/src/open_clip/model_configs/PANN-14-fmax-8k-20s.json b/src/laion_clap/clap_module/model_configs/PANN-14-fmax-8k-20s.json
similarity index 100%
rename from src/open_clip/model_configs/PANN-14-fmax-8k-20s.json
rename to src/laion_clap/clap_module/model_configs/PANN-14-fmax-8k-20s.json
diff --git a/src/open_clip/model_configs/PANN-14-tiny-transformer.json b/src/laion_clap/clap_module/model_configs/PANN-14-tiny-transformer.json
similarity index 100%
rename from src/open_clip/model_configs/PANN-14-tiny-transformer.json
rename to src/laion_clap/clap_module/model_configs/PANN-14-tiny-transformer.json
diff --git a/src/open_clip/model_configs/PANN-14-win-1536.json b/src/laion_clap/clap_module/model_configs/PANN-14-win-1536.json
similarity index 100%
rename from src/open_clip/model_configs/PANN-14-win-1536.json
rename to src/laion_clap/clap_module/model_configs/PANN-14-win-1536.json
diff --git a/src/open_clip/model_configs/PANN-14.json b/src/laion_clap/clap_module/model_configs/PANN-14.json
similarity index 100%
rename from src/open_clip/model_configs/PANN-14.json
rename to src/laion_clap/clap_module/model_configs/PANN-14.json
diff --git a/src/open_clip/model_configs/PANN-6.json b/src/laion_clap/clap_module/model_configs/PANN-6.json
similarity index 100%
rename from src/open_clip/model_configs/PANN-6.json
rename to src/laion_clap/clap_module/model_configs/PANN-6.json
diff --git a/src/open_clip/model_configs/RN101-quickgelu.json b/src/laion_clap/clap_module/model_configs/RN101-quickgelu.json
similarity index 100%
rename from src/open_clip/model_configs/RN101-quickgelu.json
rename to src/laion_clap/clap_module/model_configs/RN101-quickgelu.json
diff --git a/src/open_clip/model_configs/RN101.json b/src/laion_clap/clap_module/model_configs/RN101.json
similarity index 100%
rename from src/open_clip/model_configs/RN101.json
rename to src/laion_clap/clap_module/model_configs/RN101.json
diff --git a/src/open_clip/model_configs/RN50-quickgelu.json b/src/laion_clap/clap_module/model_configs/RN50-quickgelu.json
similarity index 100%
rename from src/open_clip/model_configs/RN50-quickgelu.json
rename to src/laion_clap/clap_module/model_configs/RN50-quickgelu.json
diff --git a/src/open_clip/model_configs/RN50.json b/src/laion_clap/clap_module/model_configs/RN50.json
similarity index 100%
rename from src/open_clip/model_configs/RN50.json
rename to src/laion_clap/clap_module/model_configs/RN50.json
diff --git a/src/open_clip/model_configs/RN50x16.json b/src/laion_clap/clap_module/model_configs/RN50x16.json
similarity index 100%
rename from src/open_clip/model_configs/RN50x16.json
rename to src/laion_clap/clap_module/model_configs/RN50x16.json
diff --git a/src/open_clip/model_configs/RN50x4.json b/src/laion_clap/clap_module/model_configs/RN50x4.json
similarity index 100%
rename from src/open_clip/model_configs/RN50x4.json
rename to src/laion_clap/clap_module/model_configs/RN50x4.json
diff --git a/src/open_clip/model_configs/ViT-B-16.json b/src/laion_clap/clap_module/model_configs/ViT-B-16.json
similarity index 100%
rename from src/open_clip/model_configs/ViT-B-16.json
rename to src/laion_clap/clap_module/model_configs/ViT-B-16.json
diff --git a/src/open_clip/model_configs/ViT-B-32-quickgelu.json b/src/laion_clap/clap_module/model_configs/ViT-B-32-quickgelu.json
similarity index 100%
rename from src/open_clip/model_configs/ViT-B-32-quickgelu.json
rename to src/laion_clap/clap_module/model_configs/ViT-B-32-quickgelu.json
diff --git a/src/open_clip/model_configs/ViT-B-32.json b/src/laion_clap/clap_module/model_configs/ViT-B-32.json
similarity index 100%
rename from src/open_clip/model_configs/ViT-B-32.json
rename to src/laion_clap/clap_module/model_configs/ViT-B-32.json
diff --git a/src/open_clip/model_configs/ViT-L-14.json b/src/laion_clap/clap_module/model_configs/ViT-L-14.json
similarity index 100%
rename from src/open_clip/model_configs/ViT-L-14.json
rename to src/laion_clap/clap_module/model_configs/ViT-L-14.json
diff --git a/src/open_clip/openai.py b/src/laion_clap/clap_module/openai.py
similarity index 100%
rename from src/open_clip/openai.py
rename to src/laion_clap/clap_module/openai.py
diff --git a/src/open_clip/pann_model.py b/src/laion_clap/clap_module/pann_model.py
similarity index 100%
rename from src/open_clip/pann_model.py
rename to src/laion_clap/clap_module/pann_model.py
diff --git a/src/open_clip/pretrained.py b/src/laion_clap/clap_module/pretrained.py
similarity index 100%
rename from src/open_clip/pretrained.py
rename to src/laion_clap/clap_module/pretrained.py
diff --git a/src/open_clip/timm_model.py b/src/laion_clap/clap_module/timm_model.py
similarity index 100%
rename from src/open_clip/timm_model.py
rename to src/laion_clap/clap_module/timm_model.py
diff --git a/src/open_clip/tokenizer.py b/src/laion_clap/clap_module/tokenizer.py
similarity index 100%
rename from src/open_clip/tokenizer.py
rename to src/laion_clap/clap_module/tokenizer.py
diff --git a/src/open_clip/transform.py b/src/laion_clap/clap_module/transform.py
similarity index 100%
rename from src/open_clip/transform.py
rename to src/laion_clap/clap_module/transform.py
diff --git a/src/open_clip/utils.py b/src/laion_clap/clap_module/utils.py
similarity index 93%
rename from src/open_clip/utils.py
rename to src/laion_clap/clap_module/utils.py
index ad9801c..b96d263 100644
--- a/src/open_clip/utils.py
+++ b/src/laion_clap/clap_module/utils.py
@@ -25,18 +25,37 @@
     "FSD50K": ["train", "test", "valid"],
     "fsd50k_class_label": ["train", "test", "valid"],
     "esc50": ["train", "test"],
+    "ESC50_1": ["train", "test"],
+    "ESC50_2": ["train", "test"],
+    "ESC50_3": ["train", "test"],
+    "ESC50_4": ["train", "test"],
+    "ESC50_5": ["train", "test"],
     "audiostock": ["train", "test"],
     "freesound_no_overlap_noesc50": ["train", "test"],
     "epidemic_sound_effects": ["train", "test"],
     "VGGSound": ["train", "test"],
     "urbansound8k_class_label": ["train", "test"],
     "audioset_t5": ["balanced_train", "unbalanced_train", "eval"],
+    "audioset_t5_debiased": ["balanced_train", "unbalanced_train", "eval"],
     "epidemic_sound_effects_t5": ["train", "test"],
+    "epidemic_sound_effects_t5_debiased": ["train", "test"],
     "WavText5K": ["train", "test"],
     "esc50_no_overlap": ["train", "test"],
     "usd8k_no_overlap": ["train", "test"],
-    "fsd50k_200_class_label": ["train", "test", "valid"]
-}
+    "fsd50k_200_class_label": ["train", "test", "valid"],
+    "fma_full": ["train", "test"],
+    "Genius": ["train", "test"],
+    "Jamendo": ["train", "test"],
+    "juno": ["train", "test"],
+    "CMU_Arctic": ["train", "test"],
+    "ravdess": ["train", "test"],
+    "Europarl-st": ["train", "test"],
+    "common_voice": ["train", "test"],
+    "Jamendo_16bit": ["train", "test"],
+    "genius_16bit_128": ["train", "test"],
+    "juno_16bit": ["train", "test"],
+    "fma_full_16bit_128": ["train", "test"],
+    }
 
 
 def freeze_batch_norm_2d(module, module_match={}, name=""):
diff --git a/src/open_clip/version.py b/src/laion_clap/clap_module/version.py
similarity index 100%
rename from src/open_clip/version.py
rename to src/laion_clap/clap_module/version.py
diff --git a/src/evaluate/__init__.py b/src/laion_clap/evaluate/__init__.py
similarity index 100%
rename from src/evaluate/__init__.py
rename to src/laion_clap/evaluate/__init__.py
diff --git a/src/evaluate/eval_dcase.py b/src/laion_clap/evaluate/eval_dcase.py
similarity index 97%
rename from src/evaluate/eval_dcase.py
rename to src/laion_clap/evaluate/eval_dcase.py
index 910bfe7..7f5057d 100644
--- a/src/evaluate/eval_dcase.py
+++ b/src/laion_clap/evaluate/eval_dcase.py
@@ -1,6 +1,4 @@
-import argparse
 import torch
-import torch.nn as nn
 import torch.nn.functional as F
 import torch.backends.cudnn as cudnn
 from open_clip import create_model
@@ -11,7 +9,7 @@
 from tqdm import tqdm
 import numpy as np
 import os
-from training.params import parse_args
+from laion_clap.training.params import parse_args
 
 
 def get_output_from_single_audio(audio, text, model, device):
@@ -125,7 +123,7 @@ def get_metrics(text_to_audio_logits):
                 from open_clip import tokenize
                 text = tokenize(text)
             else:
-                from training.data import tokenizer
+                from laion_clap.training.data import tokenizer
                 text = tokenizer(text)  # 5 texts for each audio
 
             audio_features, text_features, audio_features_mlp, text_features_mlp, logit_scale_a, logit_scale_t = \
diff --git a/src/evaluate/eval_linear_probe.py b/src/laion_clap/evaluate/eval_linear_probe.py
similarity index 97%
rename from src/evaluate/eval_linear_probe.py
rename to src/laion_clap/evaluate/eval_linear_probe.py
index bf5a524..e512345 100644
--- a/src/evaluate/eval_linear_probe.py
+++ b/src/laion_clap/evaluate/eval_linear_probe.py
@@ -1,25 +1,15 @@
 '''
 Evalute the linear probe performance on different checkpoints
 '''
-from cmath import cos
-from doctest import master
-from inspect import getargs
 import logging
 import os
-from pickletools import optimize
 import random
 from datetime import datetime
-import bisect
 import copy
 import numpy as np
 import torch
 import torch.backends.cudnn as cudnn
-from torch import optim
 from torch.cuda.amp import GradScaler
-import torch.distributed as dist
-import faulthandler
-import pathlib
-import argparse
 import glob
 
 try:
@@ -37,7 +27,7 @@
 except ImportError:
     hvd = None
 
-from open_clip import create_model_and_transforms, trace_model, create_model
+from clap_module import create_model_and_transforms, trace_model, create_model
 from training.data import get_data
 from training.params import parse_args
 from training.distributed import is_master, init_distributed_device, world_info_from_env
@@ -45,9 +35,9 @@
 from training.scheduler import cosine_lr
 from training.lp_main import config_lp_optimizer
 from training.lp_train import train_one_epoch, evaluate
-from open_clip.utils import get_tar_path_from_dataset_name, dataset_split
-from open_clip.utils import load_p, load_class_label
-from open_clip.linear_probe import LinearProbe
+from clap_module.utils import get_tar_path_from_dataset_name, dataset_split
+from clap_module.utils import load_p, load_class_label
+from clap_module.linear_probe import LinearProbe
 
 def maintain_ckpts(args, startidx, all_idx_len):
     for i in reversed(range(startidx, all_idx_len)):
diff --git a/src/laion_clap/evaluate/eval_retrieval.py b/src/laion_clap/evaluate/eval_retrieval.py
new file mode 100644
index 0000000..739734c
--- /dev/null
+++ b/src/laion_clap/evaluate/eval_retrieval.py
@@ -0,0 +1,192 @@
+import os.path
+import glob
+import random
+import numpy as np
+import logging
+import wandb
+import torch
+import torch.backends.cudnn as cudnn
+from laion_clap import create_model
+from laion_clap.training.logger import setup_logging
+from laion_clap.training.data import get_data
+from laion_clap.training.train import evaluate
+from laion_clap.utils import get_tar_path_from_dataset_name, dataset_split
+from laion_clap.training.params import parse_args
+
+
+def find_params_value(file, key):
+    # find value of params in params_file
+    with open(file, 'r') as f:
+        for line in f:
+            if key + ': ' in line:
+                return line.split(': ')[1].strip()
+    return None
+
+
+if __name__ == '__main__':
+    # (yusong) repeated run might have different metric results.
+    # This is because we randomly select crop 10s for each audio.
+    args = parse_args()
+
+    if os.path.isdir(args.pretrained):
+        log_dir = os.path.dirname(args.pretrained)
+    else:
+        log_dir = os.path.dirname(os.path.dirname(args.pretrained))
+
+    args.log_level = logging.DEBUG if args.debug else logging.INFO
+    log_path = os.path.join(log_dir, 'out.log')
+    setup_logging(log_path, args.log_level)
+    params_file = os.path.join(log_dir, 'params.txt')
+
+    seed = 3407
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+
+    cudnn.benchmark = True
+    cudnn.deterministic = False
+    pretrained = 'openai'
+    amodel = find_params_value(params_file, 'amodel')
+    tmodel = find_params_value(params_file, 'tmodel')
+
+    if amodel is None or tmodel is None:
+        raise ValueError('model type not found in params file')
+
+    # set up dummy values for args
+    args.parallel_eval = False
+    args.rank = 0
+    args.local_rank = 0
+    args.world_size = 1
+    args.val_frequency = 1
+    args.epochs = 1
+    args.precision = 'fp32'
+    args.save_logs = True
+    args.wandb = True
+    args.class_index_dict = None
+
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    args.device = device
+
+    if args.remotedata:
+        for dataset_name in args.datasetnames:
+            for split in dataset_split[dataset_name]:
+                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
+                    os.makedirs(f"./json_files/{dataset_name}/{split}")
+                os.system(
+                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
+                )
+
+    if args.datasetinfos is None:
+        args.datasetinfos = ["train", "unbalanced_train", "balanced_train"]
+    if args.dataset_type == "webdataset":
+        args.train_data = get_tar_path_from_dataset_name(
+            args.datasetnames,
+            args.datasetinfos,
+            islocal=not args.remotedata,
+            proportion=args.dataset_proportion,
+            dataset_path=args.datasetpath,
+        )
+        args.val_data = get_tar_path_from_dataset_name(
+            args.datasetnames,
+            ["valid", "test", "eval"],
+            islocal=not args.remotedata,
+            proportion=1,
+            dataset_path=args.datasetpath,
+        )
+    model, model_cfg = create_model(
+        amodel,
+        tmodel,
+        pretrained,
+        precision='fp32',
+        device=device,
+        jit=False,
+        force_quick_gelu=False,
+        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+        skip_params=False,
+        enable_fusion=args.enable_fusion,
+        fusion_type=args.fusion_type
+    )  # a hack to get model_cfg
+
+    data = get_data(args, model_cfg=model_cfg)  # (yusong): hack: no model_cfg needed to get data
+
+    writer = None  # if use tensorboard, initalize writer here
+
+    if args.wandb:
+        assert wandb is not None, "Please install wandb."
+
+        # # find the line with "wandb_notes" and get the value
+        # wandb_notes = find_params_value(params_file, 'wandb_notes')
+        # if wandb_notes is None:
+        #     print(f'wandb_notes not found in params file: {params_file}, set to timestamp.')
+        #     wandb_notes = f'experiment_{time.strftime("%Y%m%d-%H%M%S")}'
+        # wandb_notes = wandb_notes + '-eval-retrieval'
+        wandb_notes = args.wandb_notes
+
+        logging.debug("Starting wandb.")
+        args.train_sz = data["train"].dataloader.num_samples
+        if args.val_data is not None:
+            args.val_sz = data["val"].dataloader.num_samples
+        # you will have to configure this for your project!
+        if args.wandb_id is not None:
+            wandb.init(
+                project="clap",
+                id=args.wandb_id,
+                resume=True
+            )
+        else:
+            wandb.init(
+                project="clap",
+                notes=wandb_notes,
+                name=wandb_notes,
+                tags=[],
+                config=vars(args),
+            )
+        logging.debug("Finished loading wandb.")
+
+    if os.path.isdir(args.pretrained):
+        all_model_checkpoints = sorted(glob.glob(os.path.join(log_dir, 'checkpoints', '*.pt')), key=os.path.getmtime)
+    else:
+        all_model_checkpoints = [args.pretrained]
+    for model_path in all_model_checkpoints:
+        args.checkpoint_path = os.path.dirname(model_path)
+        model, model_cfg = create_model(
+            amodel,
+            tmodel,
+            pretrained,
+            precision='fp32',
+            device=device,
+            jit=False,
+            force_quick_gelu=False,
+            openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+            skip_params=False,
+            enable_fusion=args.enable_fusion,
+            fusion_type=args.fusion_type
+        )
+
+        # load model
+        checkpoint = torch.load(model_path, map_location=device)
+        if "epoch" in checkpoint:
+            # resuming a train checkpoint w/ epoch and optimizer state
+            start_epoch = checkpoint["epoch"]
+            sd = checkpoint["state_dict"]
+            if next(iter(sd.items()))[0].startswith(
+                    "module"
+            ):
+                sd = {k[len("module."):]: v for k, v in sd.items()}
+            model.load_state_dict(sd)
+            logging.info(
+                f"=> resuming checkpoint '{model_path}' (epoch {start_epoch})"
+            )
+        else:
+            # loading a bare (model only) checkpoint for fine-tune or evaluation
+            model.load_state_dict(checkpoint)
+            start_epoch = 0
+
+        model.to(device)
+        model.eval()
+        for param in model.parameters():
+            param.requires_grad = False
+
+        evaluate(model, data, start_epoch, args, writer)
diff --git a/src/evaluate/eval_retrieval_main.py b/src/laion_clap/evaluate/eval_retrieval_main.py
similarity index 90%
rename from src/evaluate/eval_retrieval_main.py
rename to src/laion_clap/evaluate/eval_retrieval_main.py
index 002f2eb..fa86c8c 100644
--- a/src/evaluate/eval_retrieval_main.py
+++ b/src/laion_clap/evaluate/eval_retrieval_main.py
@@ -1,25 +1,27 @@
-import argparse
 import os.path
 import glob
-import json
-import librosa
-from tqdm import tqdm
 import random
 import numpy as np
 import logging
 import wandb
-import time
 import torch
-import torch.nn as nn
 import torch.nn.functional as F
 import torch.backends.cudnn as cudnn
-from open_clip import create_model
-from open_clip import tokenize
+<<<<<<< HEAD:src/laion_clap/evaluate/eval_retrieval_main.py
+from clap_module import create_model
+from clap_module import tokenize
 from training.logger import setup_logging
 from training.data import get_data
 from training.train import evaluate
-from open_clip.utils import get_tar_path_from_dataset_name, dataset_split
+from clap_module.utils import get_tar_path_from_dataset_name, dataset_split
 from training.params import parse_args
+=======
+from laion_clap import create_model
+from laion_clap.training.logger import setup_logging
+from laion_clap.training.data import get_data
+from laion_clap.utils import get_tar_path_from_dataset_name, dataset_split
+from laion_clap.training.params import parse_args
+>>>>>>> 6aad20d31900bc623093e864a96f47bfd4c12cf7:src/laion_clap/evaluate/eval_zeroshot_classification.py
 
 
 def find_params_value(file, key):
@@ -55,10 +57,14 @@ def evaluate_zeroshot(model, data, start_epoch, args, writer):
         all_texts = ["This is a sound of " + t for t in args.class_index_dict.keys()]
         # (yusong): a hack, can make it better
         if args.tmodel == "transformer":
-            from open_clip import tokenize
+<<<<<<< HEAD:src/laion_clap/evaluate/eval_retrieval_main.py
+            from clap_module import tokenize
+=======
+            from laion_clap import tokenize
+>>>>>>> 6aad20d31900bc623093e864a96f47bfd4c12cf7:src/laion_clap/evaluate/eval_zeroshot_classification.py
             all_texts = tokenize(all_texts)
         else:
-            from training.data import tokenizer
+            from laion_clap.training.data import tokenizer
             all_texts = tokenizer(all_texts)
         all_text_features = model(None, all_texts, device)
         all_text_features = F.normalize(all_text_features, dim=-1).detach().cpu()
@@ -134,7 +140,7 @@ def evaluate_zeroshot(model, data, start_epoch, args, writer):
     args.epochs = 1
     args.precision = 'fp32'
     args.save_logs = True
-    args.wandb = True
+    args.wandb = args.report_to == 'wandb'
     args.class_index_dict = None
 
     device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
@@ -260,9 +266,4 @@ def evaluate_zeroshot(model, data, start_epoch, args, writer):
         for param in model.parameters():
             param.requires_grad = False
 
-        if args.datasetnames == ['esc50_no_overlap'] or args.datasetnames == ['VGGSound'] or \
-                args.datasetnames == ['usd8k_no_overlap']:
-            # use the same dataset for all models
-            evaluate_zeroshot(model, data, start_epoch, args, writer)
-        else:
-            evaluate(model, data, start_epoch, args, writer)
+        evaluate_zeroshot(model, data, start_epoch, args, writer)
diff --git a/src/laion_clap/evaluate/eval_zeroshot_classification.py b/src/laion_clap/evaluate/eval_zeroshot_classification.py
new file mode 100644
index 0000000..fa86c8c
--- /dev/null
+++ b/src/laion_clap/evaluate/eval_zeroshot_classification.py
@@ -0,0 +1,269 @@
+import os.path
+import glob
+import random
+import numpy as np
+import logging
+import wandb
+import torch
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+<<<<<<< HEAD:src/laion_clap/evaluate/eval_retrieval_main.py
+from clap_module import create_model
+from clap_module import tokenize
+from training.logger import setup_logging
+from training.data import get_data
+from training.train import evaluate
+from clap_module.utils import get_tar_path_from_dataset_name, dataset_split
+from training.params import parse_args
+=======
+from laion_clap import create_model
+from laion_clap.training.logger import setup_logging
+from laion_clap.training.data import get_data
+from laion_clap.utils import get_tar_path_from_dataset_name, dataset_split
+from laion_clap.training.params import parse_args
+>>>>>>> 6aad20d31900bc623093e864a96f47bfd4c12cf7:src/laion_clap/evaluate/eval_zeroshot_classification.py
+
+
+def find_params_value(file, key):
+    # find value of params in params_file
+    with open(file, 'r') as f:
+        for line in f:
+            if key + ': ' in line:
+                return line.split(': ')[1].strip()
+    return None
+
+
+def evaluate_zeroshot(model, data, start_epoch, args, writer):
+    dataloader = data["val"].dataloader
+    metrics = {}
+    device = torch.device(args.device)
+    model.eval()
+    metrics.update({"epoch": start_epoch})
+
+    all_audio_features = []
+    all_class_labels = []
+    with torch.no_grad():
+        for i, batch in enumerate(dataloader):
+            audios = batch  # contains mel_spec, wavform, and longer list
+            audio_features = model(audios, None, device)
+            audio_features = F.normalize(audio_features, dim=-1)
+            all_audio_features.append(audio_features.detach().cpu())
+            all_class_labels.append(torch.argmax(batch["class_label"], 1).long())
+        all_audio_features = torch.cat(all_audio_features, dim=0)
+        all_class_labels = torch.cat(all_class_labels, dim=0)
+        metrics["num_samples"] = all_audio_features.shape[0]
+
+        # get text features
+        all_texts = ["This is a sound of " + t for t in args.class_index_dict.keys()]
+        # (yusong): a hack, can make it better
+        if args.tmodel == "transformer":
+<<<<<<< HEAD:src/laion_clap/evaluate/eval_retrieval_main.py
+            from clap_module import tokenize
+=======
+            from laion_clap import tokenize
+>>>>>>> 6aad20d31900bc623093e864a96f47bfd4c12cf7:src/laion_clap/evaluate/eval_zeroshot_classification.py
+            all_texts = tokenize(all_texts)
+        else:
+            from laion_clap.training.data import tokenizer
+            all_texts = tokenizer(all_texts)
+        all_text_features = model(None, all_texts, device)
+        all_text_features = F.normalize(all_text_features, dim=-1).detach().cpu()
+
+        # compute similarity
+        logit_scale_a, logit_scale_t = model(None, None, device)
+        logit_scale_a = logit_scale_a.cpu()
+
+        logits_per_audio = (logit_scale_a * all_audio_features @ all_text_features.t()).detach().cpu()
+        logits_per_text = logits_per_audio.t().detach().cpu()
+
+        ground_truth = all_class_labels.view(-1, 1)
+        logit = logits_per_audio
+
+        ranking = torch.argsort(logit, descending=True)
+        preds = torch.where(ranking == ground_truth)[1]  # (yusong) this line is slow because it uses single thread
+        preds = preds.detach().cpu().numpy()
+        metrics[f"{args.datasetnames[0]}_mean_rank"] = preds.mean() + 1
+        metrics[f"{args.datasetnames[0]}_median_rank"] = np.floor(np.median(preds)) + 1
+        for k in [1, 5, 10]:
+            metrics[f"{args.datasetnames[0]}_R@{k}"] = np.mean(preds < k)
+        # map@10
+        metrics[f"{args.datasetnames[0]}_mAP@10"] = np.mean(np.where(preds < 10, 1 / (preds + 1), 0.0))
+
+        logging.info(
+            f"Eval Epoch: {start_epoch} "
+            + "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in metrics.items()])
+        )
+
+        if args.wandb:
+            assert wandb is not None, "Please install wandb."
+            for name, val in metrics.items():
+                wandb.log({f"val/{name}": val, "epoch": start_epoch})
+
+
+if __name__ == '__main__':
+    # (yusong) repeated run might have different metric results.
+    # This is because we randomly select crop 10s for each audio.
+    args = parse_args()
+
+    if os.path.isdir(args.pretrained):
+        log_dir = os.path.dirname(args.pretrained)
+    else:
+        log_dir = os.path.dirname(os.path.dirname(args.pretrained))
+
+    args.log_level = logging.DEBUG if args.debug else logging.INFO
+    log_path = os.path.join(log_dir, 'out.log')
+    setup_logging(log_path, args.log_level)
+    params_file = os.path.join(log_dir, 'params.txt')
+
+    seed = 3407
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+
+    cudnn.benchmark = True
+    cudnn.deterministic = False
+    pretrained = 'openai'
+    amodel = find_params_value(params_file, 'amodel')
+    tmodel = find_params_value(params_file, 'tmodel')
+
+    if amodel is None or tmodel is None:
+        raise ValueError('model type not found in params file')
+
+    # set up dummy values for args
+    args.parallel_eval = False
+    args.rank = 0
+    args.local_rank = 0
+    args.world_size = 1
+    args.val_frequency = 1
+    args.epochs = 1
+    args.precision = 'fp32'
+    args.save_logs = True
+    args.wandb = args.report_to == 'wandb'
+    args.class_index_dict = None
+
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    args.device = device
+
+    if args.remotedata:
+        for dataset_name in args.datasetnames:
+            for split in dataset_split[dataset_name]:
+                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
+                    os.makedirs(f"./json_files/{dataset_name}/{split}")
+                os.system(
+                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
+                )
+
+    if args.datasetinfos is None:
+        args.datasetinfos = ["train", "unbalanced_train", "balanced_train"]
+    if args.dataset_type == "webdataset":
+        args.train_data = get_tar_path_from_dataset_name(
+            args.datasetnames,
+            args.datasetinfos,
+            islocal=not args.remotedata,
+            proportion=args.dataset_proportion,
+            dataset_path=args.datasetpath,
+        )
+        args.val_data = get_tar_path_from_dataset_name(
+            args.datasetnames,
+            ["valid", "test", "eval"],
+            islocal=not args.remotedata,
+            proportion=1,
+            dataset_path=args.datasetpath,
+        )
+    model, model_cfg = create_model(
+        amodel,
+        tmodel,
+        pretrained,
+        precision='fp32',
+        device=device,
+        jit=False,
+        force_quick_gelu=False,
+        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+        skip_params=False,
+        enable_fusion=args.enable_fusion,
+        fusion_type=args.fusion_type
+    )  # a hack to get model_cfg
+
+    data = get_data(args, model_cfg=model_cfg)  # (yusong): hack: no model_cfg needed to get data
+
+    writer = None  # if use tensorboard, initalize writer here
+
+    if args.wandb:
+        assert wandb is not None, "Please install wandb."
+
+        # # find the line with "wandb_notes" and get the value
+        # wandb_notes = find_params_value(params_file, 'wandb_notes')
+        # if wandb_notes is None:
+        #     print(f'wandb_notes not found in params file: {params_file}, set to timestamp.')
+        #     wandb_notes = f'experiment_{time.strftime("%Y%m%d-%H%M%S")}'
+        # wandb_notes = wandb_notes + '-eval-retrieval'
+        wandb_notes = args.wandb_notes
+
+        logging.debug("Starting wandb.")
+        args.train_sz = data["train"].dataloader.num_samples
+        if args.val_data is not None:
+            args.val_sz = data["val"].dataloader.num_samples
+        # you will have to configure this for your project!
+        if args.wandb_id is not None:
+            wandb.init(
+                project="clap",
+                id=args.wandb_id,
+                resume=True
+            )
+        else:
+            wandb.init(
+                project="clap",
+                notes=wandb_notes,
+                name=wandb_notes,
+                tags=[],
+                config=vars(args),
+            )
+        logging.debug("Finished loading wandb.")
+
+    if os.path.isdir(args.pretrained):
+        all_model_checkpoints = sorted(glob.glob(os.path.join(log_dir, 'checkpoints', '*.pt')), key=os.path.getmtime)
+    else:
+        all_model_checkpoints = [args.pretrained]
+    for model_path in all_model_checkpoints:
+        args.checkpoint_path = os.path.dirname(model_path)
+        model, model_cfg = create_model(
+            amodel,
+            tmodel,
+            pretrained,
+            precision='fp32',
+            device=device,
+            jit=False,
+            force_quick_gelu=False,
+            openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+            skip_params=False,
+            enable_fusion=args.enable_fusion,
+            fusion_type=args.fusion_type
+        )
+
+        # load model
+        checkpoint = torch.load(model_path, map_location=device)
+        if "epoch" in checkpoint:
+            # resuming a train checkpoint w/ epoch and optimizer state
+            start_epoch = checkpoint["epoch"]
+            sd = checkpoint["state_dict"]
+            if next(iter(sd.items()))[0].startswith(
+                    "module"
+            ):
+                sd = {k[len("module."):]: v for k, v in sd.items()}
+            model.load_state_dict(sd)
+            logging.info(
+                f"=> resuming checkpoint '{model_path}' (epoch {start_epoch})"
+            )
+        else:
+            # loading a bare (model only) checkpoint for fine-tune or evaluation
+            model.load_state_dict(checkpoint)
+            start_epoch = 0
+
+        model.to(device)
+        model.eval()
+        for param in model.parameters():
+            param.requires_grad = False
+
+        evaluate_zeroshot(model, data, start_epoch, args, writer)
diff --git a/src/laion_clap/hook.py b/src/laion_clap/hook.py
new file mode 100644
index 0000000..ec2216b
--- /dev/null
+++ b/src/laion_clap/hook.py
@@ -0,0 +1,224 @@
+"""
+Contrastive Language-Audio Pretraining Model from LAION
+--------------------------------------------------------
+Paper: https://arxiv.org/abs/2211.06687
+Authors (equal contributions): Ke Chen, Yusong Wu, Tianyu Zhang, Yuchen Hui
+Support: LAION
+"""
+import progressbar
+import os
+import torch
+import librosa
+from clap_module import create_model
+from training.data import get_audio_features
+from training.data import int16_to_float32, float32_to_int16
+
+from transformers import RobertaTokenizer
+import wget
+from clap_module.factory import load_state_dict
+
+pbar = None
+
+def show_progress(block_num, block_size, total_size):
+    global pbar
+    if pbar is None:
+        pbar = progressbar.ProgressBar(maxval=total_size)
+        pbar.start()
+    downloaded = block_num * block_size
+    print(downloaded, total_size)
+    if downloaded < total_size:
+        pbar.update(downloaded)
+    else:
+        pbar.finish()
+        pbar = None
+
+
+class CLAP_Module(torch.nn.Module):
+    def __init__(self, enable_fusion=False, device=None, amodel= 'HTSAT-tiny', tmodel='roberta') -> None:
+        """Initialize CLAP Model
+
+        Parameters
+        ----------
+        enable_fusion: bool
+            if true, it will create the fusion clap model, otherwise non-fusion clap model (default: false) 
+        device: str
+            if None, it will automatically detect the device (gpu or cpu)
+        amodel: str
+            audio encoder architecture, default: HTSAT-tiny
+        tmodel: str
+            text encoder architecture, default: roberta
+        """
+        super(CLAP_Module, self).__init__()
+        if device is None:
+            device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+        precision = 'fp32'
+
+        if enable_fusion:
+            fusion_type = 'aff_2d'
+            model, model_cfg = create_model(
+                amodel,
+                tmodel,
+                precision=precision,
+                device=device,
+                enable_fusion=enable_fusion,
+                fusion_type=fusion_type
+            )
+        else:
+            model, model_cfg = create_model(
+                amodel,
+                tmodel,
+                precision=precision,
+                device=device,
+                enable_fusion=enable_fusion
+            )
+        self.enbale_fusion = enable_fusion
+        self.model = model
+        self.model_cfg = model_cfg
+        self.tokenize = RobertaTokenizer.from_pretrained('roberta-base')
+
+    def tokenizer(self, text):
+        result = self.tokenize(
+            text,
+            padding="max_length",
+            truncation=True,
+            max_length=77,
+            return_tensors="pt",
+        )
+        return {k: v.squeeze(0) for k, v in result.items()}
+
+    def load_ckpt(self, ckpt = None, model_id = -1):
+        """Load the pretrained checkpoint of CLAP model
+
+        Parameters
+        ----------
+        ckpt: str
+            if ckpt is specified, the model will load this ckpt, otherwise the model will download the ckpt from zenodo. \n 
+            For fusion model, it will download the 630k+audioset fusion model (id=3). For non-fusion model, it will download the 630k+audioset model (id=1).
+        model_id:
+            if model_id is specified, you can download our best ckpt, as:
+                id = 0 --> 630k non-fusion ckpt \n
+                id = 1 --> 630k+audioset non-fusion ckpt \n
+                id = 2 --> 630k fusion ckpt \n
+                id = 3 --> 630k+audioset fusion ckpt \n
+            Note that if your model is specied as non-fusion model but you download a fusion model ckpt, you will face an error.
+        """
+        download_link = 'https://huggingface.co/lukewys/laion_clap/resolve/main/'
+        download_names = [
+            '630k-best.pt',
+            '630k-audioset-best.pt',
+            '630k-fusion-best.pt',
+            '630k-audioset-fusion-best.pt'
+        ]
+        if ckpt is not None:
+            print(f'Load the specified checkpoint {ckpt} from users.')
+        else:
+            print(f'Load our best checkpoint in the paper.')
+            if model_id == -1:
+                model_id = 3 if self.enbale_fusion else 1
+            package_dir = os.path.dirname(os.path.realpath(__file__))
+            weight_file_name = download_names[model_id]
+            ckpt = os.path.join(package_dir, weight_file_name)
+            if os.path.exists(ckpt):
+                print(f'The checkpoint is already downloaded')
+            else:
+                print('Downloading laion_clap weight files...')
+                ckpt = wget.download(download_link + weight_file_name, os.path.dirname(ckpt))
+                print('Download completed!')
+        print('Load Checkpoint...')
+        ckpt = load_state_dict(ckpt, skip_params=True)
+        self.model.load_state_dict(ckpt)
+        param_names = [n for n, p in self.model.named_parameters()]
+        for n in param_names:
+            print(n, "\t", "Loaded" if n in ckpt else "Unloaded")
+    
+    def get_audio_embedding_from_filelist(self, x):
+        """get audio embeddings from the audio file list
+
+        Parameters
+        ----------
+        x: List[str] (N,): 
+            an audio file list to extract features, audio files can have different lengths (as we have the feature fusion machanism)
+        
+        Returns
+        ----------
+        audio_embed : numpy.darray (N,D):
+            audio embeddings that extracted from audio files
+        """ 
+        self.model.eval()
+        audio_input = []
+        for f in x:
+            # load the waveform of the shape (T,), should resample to 48000
+            audio_waveform, _ = librosa.load(f, sr=48000)           
+            # quantize
+            audio_waveform = int16_to_float32(float32_to_int16(audio_waveform))
+            audio_waveform = torch.from_numpy(audio_waveform).float()
+            temp_dict = {}
+            # the 'fusion' truncate mode can be changed to 'rand_trunc' if run in unfusion mode
+            temp_dict = get_audio_features(
+                temp_dict, audio_waveform, 480000, 
+                data_truncating='fusion', 
+                data_filling='repeatpad',
+                audio_cfg=self.model_cfg['audio_cfg']
+            )
+            audio_input.append(temp_dict)
+        audio_embed = self.model.get_audio_embedding(audio_input)
+        audio_embed = audio_embed.detach().cpu().numpy()
+        return audio_embed
+
+
+    def get_audio_embedding_from_data(self, x):
+        """get audio embeddings from the audio data
+
+        Parameters
+        ----------
+        x: np.darray (N,T): 
+            audio data, must be mono audio tracks.      
+        Returns
+        ----------
+        audio embed: numpy.darray (N,D):
+            audio embeddings that extracted from audio files
+        """ 
+        self.model.eval()
+        audio_input = []
+        for audio_waveform in x:          
+            # quantize
+            audio_waveform = int16_to_float32(float32_to_int16(audio_waveform))
+            audio_waveform = torch.from_numpy(audio_waveform).float()
+            temp_dict = {}
+            # the 'fusion' truncate mode can be changed to 'rand_trunc' if run in unfusion mode
+            temp_dict = get_audio_features(
+                temp_dict, audio_waveform, 480000, 
+                data_truncating='fusion', 
+                data_filling='repeatpad',
+                audio_cfg=self.model_cfg['audio_cfg']
+            )
+            audio_input.append(temp_dict)
+        audio_embed = self.model.get_audio_embedding(audio_input)
+        audio_embed = audio_embed.detach().cpu().numpy()
+        return audio_embed
+
+    def get_text_embedding(self, x, tokenizer = None):
+        """get text embeddings from texts
+
+        Parameters
+        ----------
+        x: List[str] (N,): 
+            text list 
+        tokenizer: func:
+            the tokenizer function, if not provided (None), will use the default Roberta tokenizer.
+        
+        Returns
+        ----------
+        text_embed : numpy.darray (N,D):
+            text embeddings that extracted from texts
+        """ 
+        self.model.eval()
+        if tokenizer is not None:
+            text_input = tokenizer(x)
+        else:
+            text_input = self.tokenizer(x)
+        text_embed = self.model.get_text_embedding(text_input)
+        text_embed = text_embed.detach().cpu().numpy()
+        return text_embed
+        
+    
diff --git a/src/training/__init__.py b/src/laion_clap/training/__init__.py
similarity index 100%
rename from src/training/__init__.py
rename to src/laion_clap/training/__init__.py
diff --git a/src/training/audioset_textmap.npy b/src/laion_clap/training/audioset_textmap.npy
similarity index 100%
rename from src/training/audioset_textmap.npy
rename to src/laion_clap/training/audioset_textmap.npy
diff --git a/src/training/data.py b/src/laion_clap/training/data.py
similarity index 83%
rename from src/training/data.py
rename to src/laion_clap/training/data.py
index 9e098cf..934d258 100644
--- a/src/training/data.py
+++ b/src/laion_clap/training/data.py
@@ -6,12 +6,11 @@
 import random
 import h5py
 from dataclasses import dataclass
-from training.params import parse_args
+from laion_clap.training.params import parse_args
 import braceexpand
 import numpy as np
 import pandas as pd
 import torch
-import torch.nn as nn
 import torch.nn.functional as F
 import torchvision.datasets as datasets
 import torchvision.transforms
@@ -20,13 +19,11 @@
 from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler
 from torch.utils.data.distributed import DistributedSampler
 from functools import partial
-import soundfile as sf
-import io
 from pathlib import Path
 import wget
 
-from open_clip.utils import get_tar_path_from_dataset_name, dataset_split
-from open_clip.utils import load_p, load_class_label
+from clap_module.utils import get_tar_path_from_dataset_name, dataset_split
+from clap_module.utils import load_p, load_class_label
 import tempfile
 import copy
 
@@ -42,15 +39,17 @@
 
 args = parse_args()
 if args.tmodel == "transformer":
-    from open_clip import tokenize
+    from clap_module import tokenize
 
     def tokenizer(text):
         return tokenize(text).squeeze(0)
 
 elif args.tmodel == "bert":
     from transformers import BertTokenizer
+
     tokenize = BertTokenizer.from_pretrained("bert-base-uncased")
 
+
     def tokenizer(text):
         result = tokenize(
             text,
@@ -63,8 +62,10 @@ def tokenizer(text):
 
 elif args.tmodel == "roberta":
     from transformers import RobertaTokenizer
+
     tokenize = RobertaTokenizer.from_pretrained('roberta-base')
 
+
     def tokenizer(text):
         result = tokenize(
             text,
@@ -77,8 +78,10 @@ def tokenizer(text):
 
 elif args.tmodel == "bart":
     from transformers import BartTokenizer
+
     tokenize = BartTokenizer.from_pretrained('facebook/bart-base')
 
+
     def tokenizer(text):
         result = tokenize(
             text,
@@ -103,6 +106,15 @@ def float32_to_int16(x):
     return (x * 32767.).astype(np.int16)
 
 
+def int16_to_float32_torch(x):
+    return (x / 32767.0).type(torch.float32)
+
+
+def float32_to_int16_torch(x):
+    x = torch.clamp(x, min=-1., max=1.)
+    return (x * 32767.).type(torch.int16)
+
+
 # For Toy Dataset
 class ToyDataset(Dataset):
     def __init__(self, index_path, ipc, config, eval_mode=False):
@@ -158,7 +170,7 @@ def generate_queue(self):
     def crop_wav(self, x):
         crop_size = self.audio_cfg["crop_size"]
         crop_pos = random.randint(0, len(x) - crop_size - 1)
-        return x[crop_pos : crop_pos + crop_size]
+        return x[crop_pos: crop_pos + crop_size]
 
     def prompt_text(self, target):
         events = _AUDIOSET_MAP[np.where(target > 0)]
@@ -202,10 +214,10 @@ def __getitem__(self, index):
         text = self.prompt_text(target)
         with h5py.File(hdf5_path, "r") as f:
             waveform = int16_to_float32(f["waveform"][r_idx])[
-                : self.audio_cfg["clip_samples"]
-            ]
+                       : self.audio_cfg["clip_samples"]
+                       ]
         assert (
-            len(waveform) == self.audio_cfg["clip_samples"]
+                len(waveform) == self.audio_cfg["clip_samples"]
         ), "The sample length is not match"
         # Time shift
         # if (self.config.enable_time_shift) and (not self.eval_mode):
@@ -224,11 +236,11 @@ def __getitem__(self, index):
         #             target[add_key] = 1.0
 
         # missing the text input
-        mel_spec = get_mel(torch.from_numpy(waveform), self.audio_cfg)[None,:,:]
-        mel_spec = torch.cat([mel_spec, mel_spec.clone(),mel_spec.clone(),mel_spec.clone()], dim=0).cpu().numpy()
+        mel_spec = get_mel(torch.from_numpy(waveform), self.audio_cfg)[None, :, :]
+        mel_spec = torch.cat([mel_spec, mel_spec.clone(), mel_spec.clone(), mel_spec.clone()], dim=0).cpu().numpy()
         longer = random.choice([True, False])
         if longer == False:
-            mel_spec[1:,:,:] = 0.0
+            mel_spec[1:, :, :] = 0.0
         data_dict = {
             "hdf5_path": hdf5_path,
             "index_in_hdf5": r_idx,
@@ -429,7 +441,8 @@ def sample_prop(sizefile, inputs, proportion, is_local=True):
         sampled_size_dict,
     )
 
-def get_mel(audio_data,audio_cfg):
+
+def get_mel(audio_data, audio_cfg):
     # mel shape: (n_mels, T)
     mel = torchaudio.transforms.MelSpectrogram(
         sample_rate=audio_cfg['sample_rate'],
@@ -484,7 +497,7 @@ def get_audio_features(sample, audio_data, max_len, data_truncating, data_fillin
                 # fusion
                 mel = get_mel(audio_data, audio_cfg)
                 # split to three parts
-                chunk_frames = max_len // audio_cfg['hop_size']+1  # the +1 related to how the spectrogram is computed
+                chunk_frames = max_len // audio_cfg['hop_size'] + 1  # the +1 related to how the spectrogram is computed
                 total_frames = mel.shape[0]
                 if chunk_frames == total_frames:
                     # there is a corner case where the audio length is
@@ -494,7 +507,7 @@ def get_audio_features(sample, audio_data, max_len, data_truncating, data_fillin
                     sample["mel_fusion"] = mel_fusion
                     longer = torch.tensor([False])
                 else:
-                    ranges = np.array_split(list(range(0, total_frames-chunk_frames+1)), 3)
+                    ranges = np.array_split(list(range(0, total_frames - chunk_frames + 1)), 3)
                     # print('total_frames-chunk_frames:', total_frames-chunk_frames,
                     #       'len(audio_data):', len(audio_data),
                     #       'chunk_frames:', chunk_frames,
@@ -510,16 +523,16 @@ def get_audio_features(sample, audio_data, max_len, data_truncating, data_fillin
                     idx_middle = np.random.choice(ranges[1])
                     idx_back = np.random.choice(ranges[2])
                     # select mel
-                    mel_chunk_front = mel[idx_front:idx_front+chunk_frames, :]
-                    mel_chunk_middle = mel[idx_middle:idx_middle+chunk_frames, :]
-                    mel_chunk_back = mel[idx_back:idx_back+chunk_frames, :]
+                    mel_chunk_front = mel[idx_front:idx_front + chunk_frames, :]
+                    mel_chunk_middle = mel[idx_middle:idx_middle + chunk_frames, :]
+                    mel_chunk_back = mel[idx_back:idx_back + chunk_frames, :]
 
                     # shrink the mel
                     mel_shrink = torchvision.transforms.Resize(size=[chunk_frames, 64])(mel[None])[0]
                     # logging.info(f"mel_shrink.shape: {mel_shrink.shape}")
 
                     # stack
-                    mel_fusion = torch.stack([mel_chunk_front, mel_chunk_middle, mel_chunk_back, mel_shrink], dim=0)
+                    mel_fusion = torch.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back], dim=0)
                     sample["mel_fusion"] = mel_fusion
                     longer = torch.tensor([True])
             else:
@@ -534,7 +547,7 @@ def get_audio_features(sample, audio_data, max_len, data_truncating, data_fillin
         else:  # padding if too short
             if len(audio_data) < max_len:  # do nothing if equal
                 if data_filling == "repeatpad":
-                    n_repeat = int(max_len/len(audio_data))
+                    n_repeat = int(max_len / len(audio_data))
                     audio_data = audio_data.repeat(n_repeat)
                     # audio_data = audio_data.unsqueeze(0).unsqueeze(0).unsqueeze(0)
                     # audio_data = F.interpolate(audio_data,size=max_len,mode="bicubic")[0,0,0]
@@ -552,8 +565,8 @@ def get_audio_features(sample, audio_data, max_len, data_truncating, data_fillin
                         value=0,
                     )
                 elif data_filling == "repeat":
-                    n_repeat = int(max_len/len(audio_data))
-                    audio_data = audio_data.repeat(n_repeat+1)[:max_len]
+                    n_repeat = int(max_len / len(audio_data))
+                    audio_data = audio_data.repeat(n_repeat + 1)[:max_len]
                 else:
                     raise NotImplementedError(
                         f"data_filling {data_filling} not implemented"
@@ -570,47 +583,7 @@ def get_audio_features(sample, audio_data, max_len, data_truncating, data_fillin
     return sample
 
 
-def preprocess(
-    sample,
-    audio_ext,
-    text_ext,
-    max_len,
-    audio_cfg,
-    class_index_dict=None,
-    data_filling="pad",
-    data_truncating="rand_trunc",
-    text_augment_selection=None,
-):
-    """
-    Preprocess a single sample for wdsdataloader.
-    """
-    audio_data, orig_sr = sf.read(io.BytesIO(sample[audio_ext]))
-    audio_data = int16_to_float32(float32_to_int16(audio_data))
-    audio_data = torch.tensor(audio_data).float()
-
-    # TODO: (yusong) to be include in the future
-    # # if torchaudio not installed, use soundfile to load audio
-    # if torchaudio is None:
-    #     audio_data, orig_sr = sf.read(io.BytesIO(sample[audio_ext]))
-    #     audio_data = torch.tensor(audio_data).float()
-    # else:
-    #     # https://github.com/webdataset/webdataset/blob/main/webdataset/autodecode.py
-    #     with tempfile.TemporaryDirectory() as dirname:
-    #         os.makedirs(dirname, exist_ok=True)
-    #         fname = os.path.join(dirname, f"file.flac")
-    #         with open(fname, "wb") as stream:
-    #             stream.write(sample[audio_ext])
-    #         audio_data, orig_sr = torchaudio.load(fname)
-    #         audio_data = audio_data[0, :].float()
-
-    sample = get_audio_features(sample, audio_data, max_len, data_truncating, data_filling, audio_cfg)
-    del sample[audio_ext]
-
-    try:
-        json_dict_raw = json.loads(sample[text_ext].decode("utf-8"))
-    except:
-        print("sample[__url__]:", sample["__url__"])
-
+def select_text(json_dict_raw, text_augment_selection):
     # For selecting augmented text from dataset
     if text_augment_selection is None or text_augment_selection == "none":
         texts = json_dict_raw["text"]
@@ -631,6 +604,32 @@ def preprocess(
         raise NotImplementedError(
             f"text_augment_selection {text_augment_selection} not implemented"
         )
+    return texts
+
+
+def preprocess_single(
+        sample,
+        audio_ext,
+        text_ext,
+        max_len,
+        audio_cfg,
+        class_index_dict,
+        data_filling,
+        data_truncating,
+        text_augment_selection,
+):
+    """
+    Preprocess a single sample for wdsdataloader.
+    """
+    audio_data, orig_sr = sample[audio_ext]
+    audio_data = int16_to_float32_torch(float32_to_int16_torch(audio_data[0]))
+
+    sample = get_audio_features(sample, audio_data, max_len, data_truncating, data_filling, audio_cfg)
+    del sample[audio_ext]
+
+    json_dict_raw = sample[text_ext]
+
+    texts = select_text(json_dict_raw, text_augment_selection)
     sample["full_text"] = texts
 
     if isinstance(texts, list) and isinstance(texts[0], str) and len(texts) > 1:
@@ -640,13 +639,17 @@ def preprocess(
     if class_index_dict is not None:
         # https://stackoverflow.com/questions/48004243/how-to-share-large-read-only-dictionary-list-across-processes-in-multiprocessing
         # https://stackoverflow.com/questions/45693949/storing-strings-in-a-multiprocessing-sharedctypes-array
-        # key, val = class_index_dict
-        # key = key[:].split('\n')
-        # _dict = {k: v for k, v in zip(key, val)}
-        sample["class_label"] = np.zeros(len(class_index_dict.keys()))
-        for x in json_dict_raw["tag"]:
-            sample["class_label"][class_index_dict[x]] = 1
-        sample["class_label"] = torch.tensor(sample["class_label"]).float()
+
+        # in case the re-written version is wrong, here is the old version:
+        # sample["class_label"] = np.zeros(len(class_index_dict.keys()))
+        # for x in json_dict_raw["tag"]:
+        #     sample["class_label"][class_index_dict[x]] = 1
+        # sample["class_label"] = torch.tensor(sample["class_label"]).float()
+
+        class_labels = np.zeros(len(class_index_dict))
+        class_labels[np.in1d(list(class_index_dict.keys()), json_dict_raw["tag"])] = 1
+        sample["class_label"] = torch.tensor(class_labels).float()
+
     del sample[text_ext]
     sample["audio_name"] = sample["__key__"].split("/")[-1] + "." + audio_ext
     sample["text_name"] = sample["__key__"].split("/")[-1] + "." + text_ext
@@ -654,40 +657,57 @@ def preprocess(
     return sample
 
 
-def collate_fn(batch):
+def collate_fn_with_preprocess(batch,
+                               audio_ext,
+                               text_ext,
+                               max_len,
+                               audio_cfg,
+                               class_index_dict,
+                               data_filling,
+                               data_truncating,
+                               text_augment_selection,
+                               ):
     """
     Collate function for wdsdataloader.
     batch: a list of dict, each dict is a sample
     """
     # concatenate values in each dictionary. if it is a tensor, concatenate. if it is a list, extend.
+    data_preprocessed = []
+
+    for sample in batch:
+        data_preprocessed.append(
+            preprocess_single(sample, audio_ext, text_ext, max_len, audio_cfg, class_index_dict, data_filling,
+                              data_truncating, text_augment_selection))
+
     batch_dict = {}
-    for k in batch[0].keys():
-        if isinstance(batch[0][k], dict):  # dealwith bert tokenizer output
+    for k in data_preprocessed[0].keys():
+        if isinstance(data_preprocessed[0][k], dict):  # dealwith bert tokenizer output
             batch_dict[k] = {}
-            for kk in batch[0][k].keys():
+            for kk in data_preprocessed[0][k].keys():
                 tmp = []
-                for i in range(len(batch)):
-                    tmp.append(batch[i][k][kk])
+                for i in range(len(data_preprocessed)):
+                    tmp.append(data_preprocessed[i][k][kk])
                 batch_dict[k][kk] = torch.vstack(tmp)
-        elif isinstance(batch[0][k], torch.Tensor):
-            batch_dict[k] = torch.stack([sample[k] for sample in batch])
-        elif isinstance(batch[0][k], np.ndarray):
-            batch_dict[k] = torch.tensor(np.stack([sample[k] for sample in batch]))
+        elif isinstance(data_preprocessed[0][k], torch.Tensor):
+            batch_dict[k] = torch.stack([sample[k] for sample in data_preprocessed])
+        elif isinstance(data_preprocessed[0][k], np.ndarray):
+            batch_dict[k] = torch.tensor(np.stack([sample[k] for sample in data_preprocessed]))
         else:
-            batch_dict[k] = [sample[k] for sample in batch]
+            batch_dict[k] = [sample[k] for sample in data_preprocessed]
+    del data_preprocessed
     return batch_dict
 
 
 def get_wds_dataset(
-    args,
-    model_cfg,
-    is_train,
-    audio_ext="flac",
-    text_ext="json",
-    max_len=480000,
-    proportion=1.0,
-    sizefilepath_=None,
-    is_local=None,
+        args,
+        model_cfg,
+        is_train,
+        audio_ext="flac",
+        text_ext="json",
+        max_len=480000,
+        proportion=1.0,
+        sizefilepath_=None,
+        is_local=None,
 ):
     """
     Get a dataset for wdsdataloader.
@@ -722,7 +742,7 @@ def get_wds_dataset(
                 )
         else:
             num_samples = (
-                args.val_num_samples or 0
+                    args.val_num_samples or 0
             )  # eval will just exhaust the iterator if not specified
 
     pipeline = [wds.SimpleShardList(input_shards)]
@@ -756,27 +776,42 @@ def get_wds_dataset(
                 wds.tarfile_to_samples(handler=log_and_continue),
             ]
         )
+
     pipeline.append(
-        wds.map(
-            partial(
-                preprocess,
-                audio_ext=audio_ext,
-                text_ext=text_ext,
-                max_len=max_len,
-                audio_cfg=model_cfg['audio_cfg'],
-                class_index_dict=copy.deepcopy(args.class_index_dict),
-                data_filling=args.data_filling,
-                data_truncating=args.data_truncating,
-                text_augment_selection=args.text_augment_selection,
-            )
-        ),
+        wds.decode(wds.torch_audio),
     )
 
+    # pipeline.append(
+    #     wds.map(
+    #         partial(
+    #             preprocess,
+    #             audio_ext=audio_ext,
+    #             text_ext=text_ext,
+    #             max_len=max_len,
+    #             audio_cfg=model_cfg['audio_cfg'],
+    #             class_index_dict=copy.deepcopy(args.class_index_dict),
+    #             data_filling=args.data_filling,
+    #             data_truncating=args.data_truncating,
+    #             text_augment_selection=args.text_augment_selection,
+    #         )
+    #     ),
+    # )
+
     pipeline.append(
         wds.batched(
             args.batch_size,
             partial=not (is_train or args.parallel_eval),
-            collation_fn=collate_fn,
+            collation_fn=partial(collate_fn_with_preprocess,
+                                 audio_ext=audio_ext,
+                                 text_ext=text_ext,
+                                 max_len=max_len,
+                                 audio_cfg=model_cfg['audio_cfg'],
+                                 class_index_dict=copy.deepcopy(args.class_index_dict),
+                                 data_filling=args.data_filling,
+                                 data_truncating=args.data_truncating,
+                                 text_augment_selection=args.text_augment_selection,
+                                 ),
+
         )
     )
 
@@ -804,8 +839,22 @@ def get_wds_dataset(
     if args.horovod:  # multi-node training on summit
         kwargs["multiprocessing_context"] = "forkserver"
 
+    if is_train:
+        if args.prefetch_factor:
+            prefetch_factor = args.prefetch_factor
+        else:
+            prefetch_factor = max(2, args.batch_size // args.workers)
+    else:
+        prefetch_factor = 2
+
     dataloader = wds.WebLoader(
-        dataset, batch_size=None, shuffle=False, num_workers=args.workers, **kwargs
+        dataset,
+        batch_size=None,
+        shuffle=False,
+        num_workers=args.workers,
+        pin_memory=True,
+        prefetch_factor=prefetch_factor,
+        **kwargs
     )
 
     # FIXME not clear which approach is better, with_epoch before vs after dataloader?
@@ -830,17 +879,17 @@ def get_wds_dataset(
 
 
 def wds_batch_list2dict(
-    batch,
-    keys=[
-        "__url__",
-        "__key__",
-        "waveform",
-        "text",
-        "raw_text",
-        "audio_name",
-        "text_name",
-        "audio_orig_sr",
-    ],
+        batch,
+        keys=[
+            "__url__",
+            "__key__",
+            "waveform",
+            "text",
+            "raw_text",
+            "audio_name",
+            "text_name",
+            "audio_orig_sr",
+        ],
 ):
     """
     Return a dictionary of the batch, with keys as the names of the fields.
diff --git a/src/training/distributed.py b/src/laion_clap/training/distributed.py
similarity index 100%
rename from src/training/distributed.py
rename to src/laion_clap/training/distributed.py
diff --git a/src/training/imagenet_zeroshot_data.py b/src/laion_clap/training/imagenet_zeroshot_data.py
similarity index 100%
rename from src/training/imagenet_zeroshot_data.py
rename to src/laion_clap/training/imagenet_zeroshot_data.py
diff --git a/src/training/infer_demo.py b/src/laion_clap/training/infer_demo.py
similarity index 89%
rename from src/training/infer_demo.py
rename to src/laion_clap/training/infer_demo.py
index 439ff9a..b39f3f4 100644
--- a/src/training/infer_demo.py
+++ b/src/laion_clap/training/infer_demo.py
@@ -1,7 +1,6 @@
-import os
 import torch
 import librosa
-from open_clip import create_model
+from clap_module import create_model
 from training.data import get_audio_features
 from training.data import int16_to_float32, float32_to_int16
 from transformers import RobertaTokenizer
@@ -39,9 +38,11 @@ def infer_text():
     text_data = ["I love the contrastive learning", "I love the pretrain model"] 
     # tokenize for roberta, if you want to tokenize for another text encoder, please refer to data.py#L43-90 
     text_data = tokenizer(text_data)
-    
+    model.eval()
     text_embed = model.get_text_embedding(text_data)
-    print(text_embed.size())
+    text_embed = text_embed.detach().cpu().numpy()
+    print(text_embed)
+    print(text_embed.shape)
 
 def infer_audio():
     
@@ -64,7 +65,7 @@ def infer_audio():
     )
 
     # load the waveform of the shape (T,), should resample to 48000
-    audio_waveform, sr = librosa.load('/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_long.wav', sr=48000) 
+    audio_waveform, sr = librosa.load('/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_short.wav', sr=48000) 
     # quantize
     audio_waveform = int16_to_float32(float32_to_int16(audio_waveform))
     audio_waveform = torch.from_numpy(audio_waveform).float()
@@ -77,11 +78,15 @@ def infer_audio():
         data_filling='repeatpad',
         audio_cfg=model_cfg['audio_cfg']
     )
+    model.eval()
     # can send a list to the model, to process many audio tracks in one time (i.e. batch size)
     audio_embed = model.get_audio_embedding([audio_dict])
-    print(audio_embed.size())
+    audio_embed = audio_embed.detach().cpu().numpy()
+    print(audio_embed)
+    print(audio_embed.shape)
     
 
 
 if __name__ == "__main__":
     infer_text()
+    # infer_audio()
diff --git a/src/training/logger.py b/src/laion_clap/training/logger.py
similarity index 100%
rename from src/training/logger.py
rename to src/laion_clap/training/logger.py
diff --git a/src/training/lp_main.py b/src/laion_clap/training/lp_main.py
similarity index 98%
rename from src/training/lp_main.py
rename to src/laion_clap/training/lp_main.py
index 96c7ed3..1e88f73 100644
--- a/src/training/lp_main.py
+++ b/src/laion_clap/training/lp_main.py
@@ -1,20 +1,12 @@
-from cmath import cos
-from inspect import getargs
 import logging
 import os
 import random
 from datetime import datetime
-import bisect
 import copy
-from sched import scheduler
 import numpy as np
 import torch
 import torch.backends.cudnn as cudnn
-from torch import optim
 from torch.cuda.amp import GradScaler
-import faulthandler
-import pathlib
-import argparse
 import time
 
 try:
@@ -32,16 +24,16 @@
 except ImportError:
     hvd = None
 
-from open_clip import create_model_and_transforms, trace_model, create_model
+from clap_module import create_model_and_transforms, trace_model, create_model
 from training.data import get_data
 from training.params import parse_args
 from training.distributed import is_master, init_distributed_device, world_info_from_env
 from training.logger import setup_logging
 from training.scheduler import cosine_lr
 from training.lp_train import train_one_epoch, evaluate
-from open_clip.utils import get_tar_path_from_dataset_name, dataset_split, get_optimizer
-from open_clip.utils import load_p, load_class_label
-from open_clip.linear_probe import LinearProbe
+from clap_module.utils import get_tar_path_from_dataset_name, dataset_split, get_optimizer
+from clap_module.utils import load_p, load_class_label
+from clap_module.linear_probe import LinearProbe
 
 
 def maintain_ckpts(args, startidx, all_idx_len):
diff --git a/src/training/lp_train.py b/src/laion_clap/training/lp_train.py
similarity index 98%
rename from src/training/lp_train.py
rename to src/laion_clap/training/lp_train.py
index f5a977f..d686336 100644
--- a/src/training/lp_train.py
+++ b/src/laion_clap/training/lp_train.py
@@ -14,8 +14,8 @@
 except ImportError:
     wandb = None
 
-from open_clip import LPLoss, LPMetrics, lp_gather_features
-from open_clip.utils import do_mixup, get_mix_lambda
+from clap_module import LPLoss, LPMetrics, lp_gather_features
+from clap_module.utils import do_mixup, get_mix_lambda
 from .distributed import is_master
 from .zero_shot import zero_shot_eval
 
diff --git a/src/training/main.py b/src/laion_clap/training/main.py
similarity index 98%
rename from src/training/main.py
rename to src/laion_clap/training/main.py
index fefc13b..7c48b66 100644
--- a/src/training/main.py
+++ b/src/laion_clap/training/main.py
@@ -1,17 +1,12 @@
-from inspect import getargs
 import logging
 import os
 import random
 from datetime import datetime
-import bisect
 import copy
 import numpy as np
 import torch
 import torch.backends.cudnn as cudnn
-from torch import optim
 from torch.cuda.amp import GradScaler
-import faulthandler
-import pathlib
 
 try:
     import wandb
@@ -28,14 +23,14 @@
 except ImportError:
     hvd = None
 
-from open_clip import create_model_and_transforms, trace_model, create_model
+from clap_module import create_model_and_transforms, trace_model, create_model
 from training.data import get_data
 from training.distributed import is_master, init_distributed_device, world_info_from_env
 from training.logger import setup_logging
 from training.params import parse_args
 from training.scheduler import cosine_lr
 from training.train import train_one_epoch, evaluate
-from open_clip.utils import dataset_split, get_optimizer
+from clap_module.utils import dataset_split, get_optimizer
 
 
 def maintain_ckpts(args, startidx, all_idx_len):
@@ -214,8 +209,9 @@ def main():
     assert args.precision in ["amp", "fp16", "fp32"]
     if args.precision == "fp16":
         logging.warning(
-            "It is recommended to use AMP mixed-precision instead of FP16. "
-            "FP16 support needs further verification and tuning, especially for train."
+            "It is recommended to use fp32 mixed-precision instead of FP16 and AMP in this model. "
+            "They will cause NaN loss and NaN gradients. "
+            "FP16 and AMP support needs further verification and tuning, especially for train."
         )
 
     if args.horovod:
@@ -486,6 +482,7 @@ def main():
             args.val_sz = data["val"].dataloader.num_samples
         # you will have to configure this for your project!
         wandb.init(
+            entity="clap",
             project="clap",
             notes=args.wandb_notes,
             name=args.wandb_notes,
diff --git a/src/training/params.py b/src/laion_clap/training/params.py
similarity index 98%
rename from src/training/params.py
rename to src/laion_clap/training/params.py
index a527122..84cd3b4 100644
--- a/src/training/params.py
+++ b/src/laion_clap/training/params.py
@@ -549,6 +549,12 @@ def parse_args():
         default=None,
         help="For selecting levels of augmented text. Type is among ['all', 'augment_only', 'none']",
     )
+    parser.add_argument(
+        "--prefetch-factor",
+        type=int,
+        default=None,
+        help="The prefetch factor for dataloader. Larger value will use more memory and CPU but faster.",
+    )
 
     args = parser.parse_args()
 
diff --git a/src/training/scheduler.py b/src/laion_clap/training/scheduler.py
similarity index 100%
rename from src/training/scheduler.py
rename to src/laion_clap/training/scheduler.py
diff --git a/src/training/train.py b/src/laion_clap/training/train.py
similarity index 99%
rename from src/training/train.py
rename to src/laion_clap/training/train.py
index f3e41b1..0a28476 100644
--- a/src/training/train.py
+++ b/src/laion_clap/training/train.py
@@ -14,9 +14,8 @@
 except ImportError:
     wandb = None
 
-from open_clip import ClipLoss, gather_features
+from clap_module import ClipLoss, gather_features
 from .distributed import is_master
-from .zero_shot import zero_shot_eval
 
 
 class AverageMeter(object):
@@ -598,7 +597,7 @@ def evaluate_clotho_audiocaps(
 
             # each item in the list has 5 texts
             if args.tmodel == "transformer":
-                from open_clip import tokenize
+                from clap_module import tokenize
                 texts = [tokenize(t) for t in batch['full_text']]
                 texts = torch.cat(texts)
             else:
diff --git a/src/training/zero_shot.py b/src/laion_clap/training/zero_shot.py
similarity index 99%
rename from src/training/zero_shot.py
rename to src/laion_clap/training/zero_shot.py
index 436ba9c..04472c1 100644
--- a/src/training/zero_shot.py
+++ b/src/laion_clap/training/zero_shot.py
@@ -6,7 +6,7 @@
 import torch.nn.functional as F
 from tqdm import tqdm
 
-from open_clip import tokenize
+from clap_module import tokenize
 from .imagenet_zeroshot_data import imagenet_classnames, openai_imagenet_template
 
 
diff --git a/src/training/.gitignore b/src/training/.gitignore
deleted file mode 100644
index 333c1e9..0000000
--- a/src/training/.gitignore
+++ /dev/null
@@ -1,1 +0,0 @@
-logs/


Test patch:
diff --git a/src/laion_clap/unit_test.py b/src/laion_clap/unit_test.py
new file mode 100644
index 0000000..491db7d
--- /dev/null
+++ b/src/laion_clap/unit_test.py
@@ -0,0 +1,42 @@
+"""
+Contrastive Language-Audio Pretraining Model from LAION
+--------------------------------------------------------
+Paper: https://arxiv.org/abs/2211.06687
+Authors (equal contributions): Ke Chen, Yusong Wu, Tianyu Zhang, Yuchen Hui
+Support: LAION
+"""
+
+import librosa
+import laion_clap
+
+model = laion_clap.CLAP_Module(enable_fusion=True)
+model.load_ckpt()
+
+# Directly get audio embeddings from audio files
+audio_file = [
+    '/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_short.wav',
+    '/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_long.wav'
+]
+audio_embed = model.get_audio_embedding_from_filelist(x = audio_file)
+print(audio_embed)
+print(audio_embed.shape)
+
+# Get audio embeddings from audio data
+audio_data, _ = librosa.load('/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_short.wav', sr=48000) # sample rate should be 48000
+audio_data = audio_data.reshape(1, -1) # Make it (1,T) or (N,T)
+
+audio_embed = model.get_audio_embedding_from_data(x = audio_data)
+print(audio_embed)
+print(audio_embed.shape)
+
+# Get text embedings from texts:
+text_data = ["I love the contrastive learning", "I love the pretrain model"] 
+text_embed = model.get_text_embedding(text_data)
+print(text_embed)
+print(text_embed.shape)
+
+
+
+
+
+
diff --git a/src/tests/__init__.py b/src/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/check_ckpt.py b/src/tests/check_ckpt.py
similarity index 100%
rename from src/check_ckpt.py
rename to src/tests/check_ckpt.py
diff --git a/src/test_tars.py b/src/tests/check_tars.py
similarity index 99%
rename from src/test_tars.py
rename to src/tests/check_tars.py
index f1bef55..8dcf1c1 100644
--- a/src/test_tars.py
+++ b/src/tests/check_tars.py
@@ -10,7 +10,7 @@
 import traceback
 import logging
 import json
-from open_clip import tokenize
+from laion_clap import tokenize
 
 
 def parse_args():
diff --git a/src/tests/data_loader_test.py b/src/tests/data_loader_test.py
new file mode 100644
index 0000000..03be75e
--- /dev/null
+++ b/src/tests/data_loader_test.py
@@ -0,0 +1,60 @@
+from laion_clap import create_model
+from laion_clap.training.data import get_data
+from laion_clap.training import parse_args
+import torch
+import os
+from tqdm import tqdm
+from laion_clap.training.distributed import is_master, world_info_from_env
+from laion_clap.utils import dataset_split
+
+
+def run_dataloader():
+    for i, batch in enumerate(tqdm(dataloader, total=data["train"].dataloader.num_samples // args.batch_size)):
+        pass
+
+
+if __name__ == '__main__':
+
+    args = parse_args()
+    # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule?
+    args.amodel = args.amodel.replace("/", "-")
+    device = torch.device('cpu')
+
+    # discover initial world args early so we can log properly
+    args.distributed = False
+    args.local_rank, args.rank, args.world_size = world_info_from_env()
+
+    if args.remotedata and is_master(args):
+        for dataset_name in args.datasetnames:
+            for split in dataset_split[dataset_name]:
+                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
+                    os.makedirs(f"./json_files/{dataset_name}/{split}")
+                os.system(
+                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
+                )
+
+    model, model_cfg = create_model(
+        args.amodel,
+        args.tmodel,
+        args.pretrained,
+        precision=args.precision,
+        device=device,
+        jit=args.torchscript,
+        force_quick_gelu=args.force_quick_gelu,
+        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+        skip_params=True,
+        pretrained_audio=args.pretrained_audio,
+        pretrained_text=args.pretrained_text,
+        enable_fusion=args.enable_fusion,
+        fusion_type=args.fusion_type
+    )
+
+    data = get_data(args, model_cfg)
+
+    dataloader, sampler = data["train"].dataloader, data["train"].sampler
+
+    print('dataset size:', data["train"].dataloader.num_samples)
+    print('batch size:', args.batch_size)
+    print('num batches:', data["train"].dataloader.num_samples // args.batch_size)
+
+    run_dataloader()