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

Issue:
Laion clap pip



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 52d9bee..f5499c1 100644
--- a/README.md
+++ b/README.md
@@ -2,12 +2,12 @@
 
 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>
 
 
 
-The repository contains code for the following paper:
+The repository contains code for the following paper, accepted by IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP 2023:
  - [Large-scale Contrastive Language-Audio Pretraining with Feature Fusion and Keyword-to-Caption Augmentation](https://arxiv.org/abs/2211.06687)
 
 ## About this project
@@ -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,65 +82,35 @@ 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/laion_clap/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.
 The average mAP score is calculated by averaging 4 scores: A-->T mAP@10 on AudioCaps, and T-->A mAP@10 on AudioCaps, A-->T mAP@10 on Clotho, and T-->A mAP@10 on Clotho.
 
+## Reproducibility
+An example of the preprocessed Clotho dataset in webdataset format can be download [here](https://drive.google.com/drive/folders/1mU9mBOe11jTFCrQRJQsUa4S-3TlNuYoI?usp=sharing) (by downloading, you will be agreeing the license described in the [Clotho dataset](https://zenodo.org/record/3490684#.Y9ALPeyZP1w)). The audio encoder pretrained with 48kHz AudioSet can be found [here](https://drive.google.com/drive/folders/1SMQyzJvc6DwJNuhQ_WI8tlCFL5HG2vk6?usp=sharing), where `HTSAT-fullset-imagenet-map=0.467.ckpt` is the checkpoint used to initalize our HTSAT audio encoder. You should get similar result by loading from the audio encoder checkpoint and training on same dataset.
+Because most of the dataset has copyright restriction, unfortunatly we cannot directly share other preprocessed datasets. The caption generated by keyword-to-caption model for Audioset can be found [here](https://github.com/LAION-AI/audio-dataset/tree/main/laion-audio-630k#keyword-to-caption-augmentation)
+
+
 ## Citation
 If you find this project and the LAION-Audio-630K dataset useful, please cite our paper:
 ```
-@article{wu2022large,
+@inproceedings{laionclap2023,
   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},
-  journal={arXiv preprint arXiv:2211:06687},
-  year = {2022},
+  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}
+}
+@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/eval_retrieval_freesound.sh b/experiment_scripts/eval_retrieval_freesound.sh
old mode 100755
new mode 100644
index 63d2c66..d58ac34
--- a/experiment_scripts/eval_retrieval_freesound.sh
+++ b/experiment_scripts/eval_retrieval_freesound.sh
@@ -1,44 +1,26 @@
 #!/bin/bash
 #SBATCH --comment clap
-#SBATCH --partition=gpu
+#SBATCH --partition=g40423
 #SBATCH --job-name=mclap
-#SBATCH --nodes 1
-#SBATCH --ntasks-per-node 1
+#SBATCH --nodes 3
+#SBATCH --ntasks-per-node 8
 #SBATCH --cpus-per-gpu=6
-#SBATCH --gres=gpu:1
+#SBATCH --exclusive
 #SBATCH --output=%x_%j.out
-#SBATCH --exclude gpu-st-p4d-24xlarge-[23,30,31,108,115,134,135,183,185,186,187,188,275,277,290,374]
 
-
-module load intelmpi
-source /opt/intel/mpi/latest/env/vars.sh
-export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+module load openmpi
+module load cuda/11.7
 export NCCL_PROTO=simple
-export PATH=/opt/amazon/efa/bin:$PATH
-export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
-
 export FI_EFA_FORK_SAFE=1
 export FI_LOG_LEVEL=1
 export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
-
-#export NCCL_ALGO=ring
 export NCCL_DEBUG=info
-#export NCCL_DEBUG_SUBSYS=INIT,ENV,GRAPH,COLL
-
-export PYTHONFAULTHANDLER=1
-
-export CUDA_LAUNCH_BLOCKING=0
 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
 
-
-#export NCCL_P2P_DISABLE=1
-#export NCCL_IBEXT_DISABLE=1
-#export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
-
 # sent to sub script
 export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
 export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
diff --git a/experiment_scripts/finetune-esc50.sh b/experiment_scripts/finetune-esc50.sh
old mode 100755
new mode 100644
index 2fd5929..14144ab
--- a/experiment_scripts/finetune-esc50.sh
+++ b/experiment_scripts/finetune-esc50.sh
@@ -1,44 +1,26 @@
 #!/bin/bash
 #SBATCH --comment clap
-#SBATCH --partition=gpu
+#SBATCH --partition=g40423
 #SBATCH --job-name=mclap
-#SBATCH --nodes 1
-#SBATCH --ntasks-per-node 1
+#SBATCH --nodes 3
+#SBATCH --ntasks-per-node 8
 #SBATCH --cpus-per-gpu=6
-#SBATCH --gres=gpu:1
+#SBATCH --exclusive
 #SBATCH --output=%x_%j.out
-#SBATCH --exclude gpu-st-p4d-24xlarge-[23,30,31,108,115,134,135,183,185,186,187,188,275,277,374]
 
-
-module load intelmpi
-source /opt/intel/mpi/latest/env/vars.sh
-export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+module load openmpi
+module load cuda/11.7
 export NCCL_PROTO=simple
-export PATH=/opt/amazon/efa/bin:$PATH
-export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
-
 export FI_EFA_FORK_SAFE=1
 export FI_LOG_LEVEL=1
 export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
-
-#export NCCL_ALGO=ring
 export NCCL_DEBUG=info
-#export NCCL_DEBUG_SUBSYS=INIT,ENV,GRAPH,COLL
-
-export PYTHONFAULTHANDLER=1
-
-export CUDA_LAUNCH_BLOCKING=0
 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
 
-
-#export NCCL_P2P_DISABLE=1
-#export NCCL_IBEXT_DISABLE=1
-#export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
-
 # sent to sub script
 export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
 export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
diff --git a/experiment_scripts/finetune-fsd50k.sh b/experiment_scripts/finetune-fsd50k.sh
old mode 100755
new mode 100644
index 0449420..e27d915
--- a/experiment_scripts/finetune-fsd50k.sh
+++ b/experiment_scripts/finetune-fsd50k.sh
@@ -1,44 +1,26 @@
 #!/bin/bash
 #SBATCH --comment clap
-#SBATCH --partition=gpu
+#SBATCH --partition=g40423
 #SBATCH --job-name=mclap
-#SBATCH --nodes 1
-#SBATCH --ntasks-per-node 1
+#SBATCH --nodes 3
+#SBATCH --ntasks-per-node 8
 #SBATCH --cpus-per-gpu=6
-#SBATCH --gres=gpu:1
+#SBATCH --exclusive
 #SBATCH --output=%x_%j.out
-#SBATCH --exclude gpu-st-p4d-24xlarge-[23,30,31,108,115,134,135,183,185,186,187,188,275,277,374]
 
-
-module load intelmpi
-source /opt/intel/mpi/latest/env/vars.sh
-export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+module load openmpi
+module load cuda/11.7
 export NCCL_PROTO=simple
-export PATH=/opt/amazon/efa/bin:$PATH
-export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
-
 export FI_EFA_FORK_SAFE=1
 export FI_LOG_LEVEL=1
 export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
-
-#export NCCL_ALGO=ring
 export NCCL_DEBUG=info
-#export NCCL_DEBUG_SUBSYS=INIT,ENV,GRAPH,COLL
-
-export PYTHONFAULTHANDLER=1
-
-export CUDA_LAUNCH_BLOCKING=0
 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
 
-
-#export NCCL_P2P_DISABLE=1
-#export NCCL_IBEXT_DISABLE=1
-#export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
-
 # sent to sub script
 export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
 export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
diff --git a/experiment_scripts/htsat-roberta-large-dataset-fusion.sh b/experiment_scripts/htsat-roberta-large-dataset-fusion.sh
old mode 100755
new mode 100644
index b41f432..70082a4
--- a/experiment_scripts/htsat-roberta-large-dataset-fusion.sh
+++ b/experiment_scripts/htsat-roberta-large-dataset-fusion.sh
@@ -1,43 +1,26 @@
 #!/bin/bash
 #SBATCH --comment clap
-#SBATCH --partition=gpu
+#SBATCH --partition=g40423
 #SBATCH --job-name=mclap
 #SBATCH --nodes 3
 #SBATCH --ntasks-per-node 8
 #SBATCH --cpus-per-gpu=6
-#SBATCH --gres=gpu:8
+#SBATCH --exclusive
 #SBATCH --output=%x_%j.out
-#SBATCH --exclude gpu-st-p4d-24xlarge-[23,30,31,108,115,134,135,183,185,186,187,188,275,277,374]
 
-module load intelmpi
-source /opt/intel/mpi/latest/env/vars.sh
-export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+module load openmpi
+module load cuda/11.7
 export NCCL_PROTO=simple
-export PATH=/opt/amazon/efa/bin:$PATH
-export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
-
 export FI_EFA_FORK_SAFE=1
 export FI_LOG_LEVEL=1
 export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
-
-#export NCCL_ALGO=ring
 export NCCL_DEBUG=info
-#export NCCL_DEBUG_SUBSYS=INIT,ENV,GRAPH,COLL
-
-export PYTHONFAULTHANDLER=1
-
-export CUDA_LAUNCH_BLOCKING=0
 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
 
-
-#export NCCL_P2P_DISABLE=1
-#export NCCL_IBEXT_DISABLE=1
-#export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
-
 # sent to sub script
 export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
 export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
diff --git a/experiment_scripts/train-htsat-roberta.sh b/experiment_scripts/train-htsat-roberta.sh
old mode 100755
new mode 100644
index 9982991..c2da543
--- a/experiment_scripts/train-htsat-roberta.sh
+++ b/experiment_scripts/train-htsat-roberta.sh
@@ -1,43 +1,26 @@
 #!/bin/bash
 #SBATCH --comment clap
-#SBATCH --partition=gpu
+#SBATCH --partition=g40423
 #SBATCH --job-name=mclap
-#SBATCH --nodes 1
+#SBATCH --nodes 3
 #SBATCH --ntasks-per-node 8
 #SBATCH --cpus-per-gpu=6
-#SBATCH --gres=gpu:8
+#SBATCH --exclusive
 #SBATCH --output=%x_%j.out
-#SBATCH --exclude gpu-st-p4d-24xlarge-[23,30,31,108,115,134,135,183,185,186,187,188,275,277,374]
 
-module load intelmpi
-source /opt/intel/mpi/latest/env/vars.sh
-export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+module load openmpi
+module load cuda/11.7
 export NCCL_PROTO=simple
-export PATH=/opt/amazon/efa/bin:$PATH
-export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
-
 export FI_EFA_FORK_SAFE=1
 export FI_LOG_LEVEL=1
 export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
-
-#export NCCL_ALGO=ring
 export NCCL_DEBUG=info
-#export NCCL_DEBUG_SUBSYS=INIT,ENV,GRAPH,COLL
-
-export PYTHONFAULTHANDLER=1
-
-export CUDA_LAUNCH_BLOCKING=0
 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
 
-
-#export NCCL_P2P_DISABLE=1
-#export NCCL_IBEXT_DISABLE=1
-#export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
-
 # sent to sub script
 export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
 export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
diff --git a/experiment_scripts/train-pann-roberta.sh b/experiment_scripts/train-pann-roberta.sh
old mode 100755
new mode 100644
index 4ea32d3..8294157
--- a/experiment_scripts/train-pann-roberta.sh
+++ b/experiment_scripts/train-pann-roberta.sh
@@ -1,43 +1,26 @@
 #!/bin/bash
 #SBATCH --comment clap
-#SBATCH --partition=gpu
+#SBATCH --partition=g40423
 #SBATCH --job-name=mclap
-#SBATCH --nodes 1
+#SBATCH --nodes 3
 #SBATCH --ntasks-per-node 8
 #SBATCH --cpus-per-gpu=6
-#SBATCH --gres=gpu:8
+#SBATCH --exclusive
 #SBATCH --output=%x_%j.out
-#SBATCH --exclude gpu-st-p4d-24xlarge-[23,30,31,108,115,134,135,183,185,186,187,188,275,277,374]
 
-module load intelmpi
-source /opt/intel/mpi/latest/env/vars.sh
-export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+module load openmpi
+module load cuda/11.7
 export NCCL_PROTO=simple
-export PATH=/opt/amazon/efa/bin:$PATH
-export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
-
 export FI_EFA_FORK_SAFE=1
 export FI_LOG_LEVEL=1
 export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
-
-#export NCCL_ALGO=ring
 export NCCL_DEBUG=info
-#export NCCL_DEBUG_SUBSYS=INIT,ENV,GRAPH,COLL
-
-export PYTHONFAULTHANDLER=1
-
-export CUDA_LAUNCH_BLOCKING=0
 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
 
-
-#export NCCL_P2P_DISABLE=1
-#export NCCL_IBEXT_DISABLE=1
-#export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
-
 # sent to sub script
 export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
 export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
diff --git a/experiment_scripts/zeroshot_esc50.sh b/experiment_scripts/zeroshot_esc50.sh
old mode 100755
new mode 100644
index be0cfa1..cfba10f
--- a/experiment_scripts/zeroshot_esc50.sh
+++ b/experiment_scripts/zeroshot_esc50.sh
@@ -102,4 +102,4 @@ python -m evaluate.eval_zeroshot_classification \
 --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"
\ No newline at end of file
+--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/laion_clap/clap_module/htsat.py b/src/laion_clap/clap_module/htsat.py
index db96116..bb8e7cf 100644
--- a/src/laion_clap/clap_module/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/laion_clap/clap_module/model.py b/src/laion_clap/clap_module/model.py
index e35a2c7..60663ec 100644
--- a/src/laion_clap/clap_module/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/laion_clap/evaluate/eval_dcase.py b/src/laion_clap/evaluate/eval_dcase.py
index 793bf27..7f5057d 100644
--- a/src/laion_clap/evaluate/eval_dcase.py
+++ b/src/laion_clap/evaluate/eval_dcase.py
@@ -1,7 +1,8 @@
 import torch
 import torch.nn.functional as F
 import torch.backends.cudnn as cudnn
-from laion_clap import create_model
+from open_clip import create_model
+from open_clip import tokenize
 import glob
 import json
 import librosa
@@ -119,7 +120,7 @@ def get_metrics(text_to_audio_logits):
             text = json_data["text"]
 
             if args.tmodel == "transformer":
-                from laion_clap import tokenize
+                from open_clip import tokenize
                 text = tokenize(text)
             else:
                 from laion_clap.training.data import tokenizer
diff --git a/src/laion_clap/evaluate/eval_linear_probe.py b/src/laion_clap/evaluate/eval_linear_probe.py
index d81a8c2..e512345 100644
--- a/src/laion_clap/evaluate/eval_linear_probe.py
+++ b/src/laion_clap/evaluate/eval_linear_probe.py
@@ -27,16 +27,17 @@
 except ImportError:
     hvd = None
 
-from laion_clap import trace_model, create_model
-from laion_clap.training.data import get_data
-from laion_clap.training.params import parse_args
-from laion_clap.training.distributed import is_master, init_distributed_device, world_info_from_env
-from laion_clap.training.logger import setup_logging
-from laion_clap.training.lp_main import config_lp_optimizer
-from laion_clap.training.lp_train import train_one_epoch, evaluate
-from laion_clap.utils import dataset_split
-from laion_clap.utils import load_class_label
-from laion_clap.linear_probe import LinearProbe
+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_main import config_lp_optimizer
+from training.lp_train import train_one_epoch, evaluate
+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_main.py b/src/laion_clap/evaluate/eval_retrieval_main.py
new file mode 100644
index 0000000..fa86c8c
--- /dev/null
+++ b/src/laion_clap/evaluate/eval_retrieval_main.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/evaluate/eval_zeroshot_classification.py b/src/laion_clap/evaluate/eval_zeroshot_classification.py
index 39a64c0..fa86c8c 100644
--- a/src/laion_clap/evaluate/eval_zeroshot_classification.py
+++ b/src/laion_clap/evaluate/eval_zeroshot_classification.py
@@ -7,11 +7,21 @@
 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):
@@ -47,7 +57,11 @@ 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":
+<<<<<<< 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
diff --git a/src/laion_clap/hook.py b/src/laion_clap/hook.py
index bb5a33f..ec2216b 100644
--- a/src/laion_clap/hook.py
+++ b/src/laion_clap/hook.py
@@ -10,15 +10,15 @@
 import torch
 import librosa
 from clap_module import create_model
-from laion_clap.training.data import get_audio_features
-from laion_clap.training.data import int16_to_float32, float32_to_int16
+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:
@@ -34,12 +34,13 @@ def show_progress(block_num, block_size, total_size):
 
 
 class CLAP_Module(torch.nn.Module):
-    def __init__(self, enable_fusion=False, device=None, amodel='HTSAT-tiny', tmodel='roberta') -> None:
+    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)
+            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
@@ -85,12 +86,13 @@ def tokenizer(self, text):
         )
         return {k: v.squeeze(0) for k, v in result.items()}
 
-    def load_ckpt(self, ckpt=None, model_id=-1):
+    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
+            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:
@@ -128,32 +130,33 @@ def load_ckpt(self, ckpt=None, model_id=-1):
         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,):
+        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)
+            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',
+                temp_dict, audio_waveform, 480000, 
+                data_truncating='fusion', 
                 data_filling='repeatpad',
                 audio_cfg=self.model_cfg['audio_cfg']
             )
@@ -162,28 +165,30 @@ def get_audio_embedding_from_filelist(self, x):
         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.
+        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:
+        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',
+                temp_dict, audio_waveform, 480000, 
+                data_truncating='fusion', 
                 data_filling='repeatpad',
                 audio_cfg=self.model_cfg['audio_cfg']
             )
@@ -192,20 +197,21 @@ def get_audio_embedding_from_data(self, x):
         audio_embed = audio_embed.detach().cpu().numpy()
         return audio_embed
 
-    def get_text_embedding(self, x, tokenizer=None):
+    def get_text_embedding(self, x, tokenizer = None):
         """get text embeddings from texts
+
         Parameters
         ----------
-        x: List[str] (N,):
-            text list
+        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)
@@ -213,4 +219,6 @@ def get_text_embedding(self, x, tokenizer=None):
             text_input = self.tokenizer(x)
         text_embed = self.model.get_text_embedding(text_input)
         text_embed = text_embed.detach().cpu().numpy()
-        return text_embed
\ No newline at end of file
+        return text_embed
+        
+    
diff --git a/src/laion_clap/training/data.py b/src/laion_clap/training/data.py
index b152b3a..934d258 100644
--- a/src/laion_clap/training/data.py
+++ b/src/laion_clap/training/data.py
@@ -22,8 +22,9 @@
 from pathlib import Path
 import wget
 
-from laion_clap.utils import get_tar_path_from_dataset_name, dataset_split
-from laion_clap.utils import 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
 
 try:
@@ -38,8 +39,7 @@
 
 args = parse_args()
 if args.tmodel == "transformer":
-    from laion_clap import tokenize
-
+    from clap_module import tokenize
 
     def tokenizer(text):
         return tokenize(text).squeeze(0)
diff --git a/src/laion_clap/training/imagenet_zeroshot_data.py b/src/laion_clap/training/imagenet_zeroshot_data.py
new file mode 100644
index 0000000..a789874
--- /dev/null
+++ b/src/laion_clap/training/imagenet_zeroshot_data.py
@@ -0,0 +1,254 @@
+# NOTE: This script is currently not supported for CLAP.
+
+imagenet_classnames = ["tench", "goldfish", "great white shark", "tiger shark", "hammerhead shark", "electric ray",
+                        "stingray", "rooster", "hen", "ostrich", "brambling", "goldfinch", "house finch", "junco",
+                        "indigo bunting", "American robin", "bulbul", "jay", "magpie", "chickadee", "American dipper",
+                        "kite (bird of prey)", "bald eagle", "vulture", "great grey owl", "fire salamander",
+                        "smooth newt", "newt", "spotted salamander", "axolotl", "American bullfrog", "tree frog",
+                        "tailed frog", "loggerhead sea turtle", "leatherback sea turtle", "mud turtle", "terrapin",
+                        "box turtle", "banded gecko", "green iguana", "Carolina anole",
+                        "desert grassland whiptail lizard", "agama", "frilled-necked lizard", "alligator lizard",
+                        "Gila monster", "European green lizard", "chameleon", "Komodo dragon", "Nile crocodile",
+                        "American alligator", "triceratops", "worm snake", "ring-necked snake",
+                        "eastern hog-nosed snake", "smooth green snake", "kingsnake", "garter snake", "water snake",
+                        "vine snake", "night snake", "boa constrictor", "African rock python", "Indian cobra",
+                        "green mamba", "sea snake", "Saharan horned viper", "eastern diamondback rattlesnake",
+                        "sidewinder rattlesnake", "trilobite", "harvestman", "scorpion", "yellow garden spider",
+                        "barn spider", "European garden spider", "southern black widow", "tarantula", "wolf spider",
+                        "tick", "centipede", "black grouse", "ptarmigan", "ruffed grouse", "prairie grouse", "peafowl",
+                        "quail", "partridge", "african grey parrot", "macaw", "sulphur-crested cockatoo", "lorikeet",
+                        "coucal", "bee eater", "hornbill", "hummingbird", "jacamar", "toucan", "duck",
+                        "red-breasted merganser", "goose", "black swan", "tusker", "echidna", "platypus", "wallaby",
+                        "koala", "wombat", "jellyfish", "sea anemone", "brain coral", "flatworm", "nematode", "conch",
+                        "snail", "slug", "sea slug", "chiton", "chambered nautilus", "Dungeness crab", "rock crab",
+                        "fiddler crab", "red king crab", "American lobster", "spiny lobster", "crayfish", "hermit crab",
+                        "isopod", "white stork", "black stork", "spoonbill", "flamingo", "little blue heron",
+                        "great egret", "bittern bird", "crane bird", "limpkin", "common gallinule", "American coot",
+                        "bustard", "ruddy turnstone", "dunlin", "common redshank", "dowitcher", "oystercatcher",
+                        "pelican", "king penguin", "albatross", "grey whale", "killer whale", "dugong", "sea lion",
+                        "Chihuahua", "Japanese Chin", "Maltese", "Pekingese", "Shih Tzu", "King Charles Spaniel",
+                        "Papillon", "toy terrier", "Rhodesian Ridgeback", "Afghan Hound", "Basset Hound", "Beagle",
+                        "Bloodhound", "Bluetick Coonhound", "Black and Tan Coonhound", "Treeing Walker Coonhound",
+                        "English foxhound", "Redbone Coonhound", "borzoi", "Irish Wolfhound", "Italian Greyhound",
+                        "Whippet", "Ibizan Hound", "Norwegian Elkhound", "Otterhound", "Saluki", "Scottish Deerhound",
+                        "Weimaraner", "Staffordshire Bull Terrier", "American Staffordshire Terrier",
+                        "Bedlington Terrier", "Border Terrier", "Kerry Blue Terrier", "Irish Terrier",
+                        "Norfolk Terrier", "Norwich Terrier", "Yorkshire Terrier", "Wire Fox Terrier",
+                        "Lakeland Terrier", "Sealyham Terrier", "Airedale Terrier", "Cairn Terrier",
+                        "Australian Terrier", "Dandie Dinmont Terrier", "Boston Terrier", "Miniature Schnauzer",
+                        "Giant Schnauzer", "Standard Schnauzer", "Scottish Terrier", "Tibetan Terrier",
+                        "Australian Silky Terrier", "Soft-coated Wheaten Terrier", "West Highland White Terrier",
+                        "Lhasa Apso", "Flat-Coated Retriever", "Curly-coated Retriever", "Golden Retriever",
+                        "Labrador Retriever", "Chesapeake Bay Retriever", "German Shorthaired Pointer", "Vizsla",
+                        "English Setter", "Irish Setter", "Gordon Setter", "Brittany dog", "Clumber Spaniel",
+                        "English Springer Spaniel", "Welsh Springer Spaniel", "Cocker Spaniel", "Sussex Spaniel",
+                        "Irish Water Spaniel", "Kuvasz", "Schipperke", "Groenendael dog", "Malinois", "Briard",
+                        "Australian Kelpie", "Komondor", "Old English Sheepdog", "Shetland Sheepdog", "collie",
+                        "Border Collie", "Bouvier des Flandres dog", "Rottweiler", "German Shepherd Dog", "Dobermann",
+                        "Miniature Pinscher", "Greater Swiss Mountain Dog", "Bernese Mountain Dog",
+                        "Appenzeller Sennenhund", "Entlebucher Sennenhund", "Boxer", "Bullmastiff", "Tibetan Mastiff",
+                        "French Bulldog", "Great Dane", "St. Bernard", "husky", "Alaskan Malamute", "Siberian Husky",
+                        "Dalmatian", "Affenpinscher", "Basenji", "pug", "Leonberger", "Newfoundland dog",
+                        "Great Pyrenees dog", "Samoyed", "Pomeranian", "Chow Chow", "Keeshond", "brussels griffon",
+                        "Pembroke Welsh Corgi", "Cardigan Welsh Corgi", "Toy Poodle", "Miniature Poodle",
+                        "Standard Poodle", "Mexican hairless dog (xoloitzcuintli)", "grey wolf", "Alaskan tundra wolf",
+                        "red wolf or maned wolf", "coyote", "dingo", "dhole", "African wild dog", "hyena", "red fox",
+                        "kit fox", "Arctic fox", "grey fox", "tabby cat", "tiger cat", "Persian cat", "Siamese cat",
+                        "Egyptian Mau", "cougar", "lynx", "leopard", "snow leopard", "jaguar", "lion", "tiger",
+                        "cheetah", "brown bear", "American black bear", "polar bear", "sloth bear", "mongoose",
+                        "meerkat", "tiger beetle", "ladybug", "ground beetle", "longhorn beetle", "leaf beetle",
+                        "dung beetle", "rhinoceros beetle", "weevil", "fly", "bee", "ant", "grasshopper",
+                        "cricket insect", "stick insect", "cockroach", "praying mantis", "cicada", "leafhopper",
+                        "lacewing", "dragonfly", "damselfly", "red admiral butterfly", "ringlet butterfly",
+                        "monarch butterfly", "small white butterfly", "sulphur butterfly", "gossamer-winged butterfly",
+                        "starfish", "sea urchin", "sea cucumber", "cottontail rabbit", "hare", "Angora rabbit",
+                        "hamster", "porcupine", "fox squirrel", "marmot", "beaver", "guinea pig", "common sorrel horse",
+                        "zebra", "pig", "wild boar", "warthog", "hippopotamus", "ox", "water buffalo", "bison",
+                        "ram (adult male sheep)", "bighorn sheep", "Alpine ibex", "hartebeest", "impala (antelope)",
+                        "gazelle", "arabian camel", "llama", "weasel", "mink", "European polecat",
+                        "black-footed ferret", "otter", "skunk", "badger", "armadillo", "three-toed sloth", "orangutan",
+                        "gorilla", "chimpanzee", "gibbon", "siamang", "guenon", "patas monkey", "baboon", "macaque",
+                        "langur", "black-and-white colobus", "proboscis monkey", "marmoset", "white-headed capuchin",
+                        "howler monkey", "titi monkey", "Geoffroy's spider monkey", "common squirrel monkey",
+                        "ring-tailed lemur", "indri", "Asian elephant", "African bush elephant", "red panda",
+                        "giant panda", "snoek fish", "eel", "silver salmon", "rock beauty fish", "clownfish",
+                        "sturgeon", "gar fish", "lionfish", "pufferfish", "abacus", "abaya", "academic gown",
+                        "accordion", "acoustic guitar", "aircraft carrier", "airliner", "airship", "altar", "ambulance",
+                        "amphibious vehicle", "analog clock", "apiary", "apron", "trash can", "assault rifle",
+                        "backpack", "bakery", "balance beam", "balloon", "ballpoint pen", "Band-Aid", "banjo",
+                        "baluster / handrail", "barbell", "barber chair", "barbershop", "barn", "barometer", "barrel",
+                        "wheelbarrow", "baseball", "basketball", "bassinet", "bassoon", "swimming cap", "bath towel",
+                        "bathtub", "station wagon", "lighthouse", "beaker", "military hat (bearskin or shako)",
+                        "beer bottle", "beer glass", "bell tower", "baby bib", "tandem bicycle", "bikini",
+                        "ring binder", "binoculars", "birdhouse", "boathouse", "bobsleigh", "bolo tie", "poke bonnet",
+                        "bookcase", "bookstore", "bottle cap", "hunting bow", "bow tie", "brass memorial plaque", "bra",
+                        "breakwater", "breastplate", "broom", "bucket", "buckle", "bulletproof vest",
+                        "high-speed train", "butcher shop", "taxicab", "cauldron", "candle", "cannon", "canoe",
+                        "can opener", "cardigan", "car mirror", "carousel", "tool kit", "cardboard box / carton",
+                        "car wheel", "automated teller machine", "cassette", "cassette player", "castle", "catamaran",
+                        "CD player", "cello", "mobile phone", "chain", "chain-link fence", "chain mail", "chainsaw",
+                        "storage chest", "chiffonier", "bell or wind chime", "china cabinet", "Christmas stocking",
+                        "church", "movie theater", "cleaver", "cliff dwelling", "cloak", "clogs", "cocktail shaker",
+                        "coffee mug", "coffeemaker", "spiral or coil", "combination lock", "computer keyboard",
+                        "candy store", "container ship", "convertible", "corkscrew", "cornet", "cowboy boot",
+                        "cowboy hat", "cradle", "construction crane", "crash helmet", "crate", "infant bed",
+                        "Crock Pot", "croquet ball", "crutch", "cuirass", "dam", "desk", "desktop computer",
+                        "rotary dial telephone", "diaper", "digital clock", "digital watch", "dining table",
+                        "dishcloth", "dishwasher", "disc brake", "dock", "dog sled", "dome", "doormat", "drilling rig",
+                        "drum", "drumstick", "dumbbell", "Dutch oven", "electric fan", "electric guitar",
+                        "electric locomotive", "entertainment center", "envelope", "espresso machine", "face powder",
+                        "feather boa", "filing cabinet", "fireboat", "fire truck", "fire screen", "flagpole", "flute",
+                        "folding chair", "football helmet", "forklift", "fountain", "fountain pen", "four-poster bed",
+                        "freight car", "French horn", "frying pan", "fur coat", "garbage truck",
+                        "gas mask or respirator", "gas pump", "goblet", "go-kart", "golf ball", "golf cart", "gondola",
+                        "gong", "gown", "grand piano", "greenhouse", "radiator grille", "grocery store", "guillotine",
+                        "hair clip", "hair spray", "half-track", "hammer", "hamper", "hair dryer", "hand-held computer",
+                        "handkerchief", "hard disk drive", "harmonica", "harp", "combine harvester", "hatchet",
+                        "holster", "home theater", "honeycomb", "hook", "hoop skirt", "gymnastic horizontal bar",
+                        "horse-drawn vehicle", "hourglass", "iPod", "clothes iron", "carved pumpkin", "jeans", "jeep",
+                        "T-shirt", "jigsaw puzzle", "rickshaw", "joystick", "kimono", "knee pad", "knot", "lab coat",
+                        "ladle", "lampshade", "laptop computer", "lawn mower", "lens cap", "letter opener", "library",
+                        "lifeboat", "lighter", "limousine", "ocean liner", "lipstick", "slip-on shoe", "lotion",
+                        "music speaker", "loupe magnifying glass", "sawmill", "magnetic compass", "messenger bag",
+                        "mailbox", "tights", "one-piece bathing suit", "manhole cover", "maraca", "marimba", "mask",
+                        "matchstick", "maypole", "maze", "measuring cup", "medicine cabinet", "megalith", "microphone",
+                        "microwave oven", "military uniform", "milk can", "minibus", "miniskirt", "minivan", "missile",
+                        "mitten", "mixing bowl", "mobile home", "ford model t", "modem", "monastery", "monitor",
+                        "moped", "mortar and pestle", "graduation cap", "mosque", "mosquito net", "vespa",
+                        "mountain bike", "tent", "computer mouse", "mousetrap", "moving van", "muzzle", "metal nail",
+                        "neck brace", "necklace", "baby pacifier", "notebook computer", "obelisk", "oboe", "ocarina",
+                        "odometer", "oil filter", "pipe organ", "oscilloscope", "overskirt", "bullock cart",
+                        "oxygen mask", "product packet / packaging", "paddle", "paddle wheel", "padlock", "paintbrush",
+                        "pajamas", "palace", "pan flute", "paper towel", "parachute", "parallel bars", "park bench",
+                        "parking meter", "railroad car", "patio", "payphone", "pedestal", "pencil case",
+                        "pencil sharpener", "perfume", "Petri dish", "photocopier", "plectrum", "Pickelhaube",
+                        "picket fence", "pickup truck", "pier", "piggy bank", "pill bottle", "pillow", "ping-pong ball",
+                        "pinwheel", "pirate ship", "drink pitcher", "block plane", "planetarium", "plastic bag",
+                        "plate rack", "farm plow", "plunger", "Polaroid camera", "pole", "police van", "poncho",
+                        "pool table", "soda bottle", "plant pot", "potter's wheel", "power drill", "prayer rug",
+                        "printer", "prison", "missile", "projector", "hockey puck", "punching bag", "purse", "quill",
+                        "quilt", "race car", "racket", "radiator", "radio", "radio telescope", "rain barrel",
+                        "recreational vehicle", "fishing casting reel", "reflex camera", "refrigerator",
+                        "remote control", "restaurant", "revolver", "rifle", "rocking chair", "rotisserie", "eraser",
+                        "rugby ball", "ruler measuring stick", "sneaker", "safe", "safety pin", "salt shaker", "sandal",
+                        "sarong", "saxophone", "scabbard", "weighing scale", "school bus", "schooner", "scoreboard",
+                        "CRT monitor", "screw", "screwdriver", "seat belt", "sewing machine", "shield", "shoe store",
+                        "shoji screen / room divider", "shopping basket", "shopping cart", "shovel", "shower cap",
+                        "shower curtain", "ski", "balaclava ski mask", "sleeping bag", "slide rule", "sliding door",
+                        "slot machine", "snorkel", "snowmobile", "snowplow", "soap dispenser", "soccer ball", "sock",
+                        "solar thermal collector", "sombrero", "soup bowl", "keyboard space bar", "space heater",
+                        "space shuttle", "spatula", "motorboat", "spider web", "spindle", "sports car", "spotlight",
+                        "stage", "steam locomotive", "through arch bridge", "steel drum", "stethoscope", "scarf",
+                        "stone wall", "stopwatch", "stove", "strainer", "tram", "stretcher", "couch", "stupa",
+                        "submarine", "suit", "sundial", "sunglasses", "sunglasses", "sunscreen", "suspension bridge",
+                        "mop", "sweatshirt", "swim trunks / shorts", "swing", "electrical switch", "syringe",
+                        "table lamp", "tank", "tape player", "teapot", "teddy bear", "television", "tennis ball",
+                        "thatched roof", "front curtain", "thimble", "threshing machine", "throne", "tile roof",
+                        "toaster", "tobacco shop", "toilet seat", "torch", "totem pole", "tow truck", "toy store",
+                        "tractor", "semi-trailer truck", "tray", "trench coat", "tricycle", "trimaran", "tripod",
+                        "triumphal arch", "trolleybus", "trombone", "hot tub", "turnstile", "typewriter keyboard",
+                        "umbrella", "unicycle", "upright piano", "vacuum cleaner", "vase", "vaulted or arched ceiling",
+                        "velvet fabric", "vending machine", "vestment", "viaduct", "violin", "volleyball",
+                        "waffle iron", "wall clock", "wallet", "wardrobe", "military aircraft", "sink",
+                        "washing machine", "water bottle", "water jug", "water tower", "whiskey jug", "whistle",
+                        "hair wig", "window screen", "window shade", "Windsor tie", "wine bottle", "airplane wing",
+                        "wok", "wooden spoon", "wool", "split-rail fence", "shipwreck", "sailboat", "yurt", "website",
+                        "comic book", "crossword", "traffic or street sign", "traffic light", "dust jacket", "menu",
+                        "plate", "guacamole", "consomme", "hot pot", "trifle", "ice cream", "popsicle", "baguette",
+                        "bagel", "pretzel", "cheeseburger", "hot dog", "mashed potatoes", "cabbage", "broccoli",
+                        "cauliflower", "zucchini", "spaghetti squash", "acorn squash", "butternut squash", "cucumber",
+                        "artichoke", "bell pepper", "cardoon", "mushroom", "Granny Smith apple", "strawberry", "orange",
+                        "lemon", "fig", "pineapple", "banana", "jackfruit", "cherimoya (custard apple)", "pomegranate",
+                        "hay", "carbonara", "chocolate syrup", "dough", "meatloaf", "pizza", "pot pie", "burrito",
+                        "red wine", "espresso", "tea cup", "eggnog", "mountain", "bubble", "cliff", "coral reef",
+                        "geyser", "lakeshore", "promontory", "sandbar", "beach", "valley", "volcano", "baseball player",
+                        "bridegroom", "scuba diver", "rapeseed", "daisy", "yellow lady's slipper", "corn", "acorn",
+                        "rose hip", "horse chestnut seed", "coral fungus", "agaric", "gyromitra", "stinkhorn mushroom",
+                        "earth star fungus", "hen of the woods mushroom", "bolete", "corn cob", "toilet paper"]
+
+
+
+
+
+openai_imagenet_template = [
+    lambda c: f'a bad photo of a {c}.',
+    lambda c: f'a photo of many {c}.',
+    lambda c: f'a sculpture of a {c}.',
+    lambda c: f'a photo of the hard to see {c}.',
+    lambda c: f'a low resolution photo of the {c}.',
+    lambda c: f'a rendering of a {c}.',
+    lambda c: f'graffiti of a {c}.',
+    lambda c: f'a bad photo of the {c}.',
+    lambda c: f'a cropped photo of the {c}.',
+    lambda c: f'a tattoo of a {c}.',
+    lambda c: f'the embroidered {c}.',
+    lambda c: f'a photo of a hard to see {c}.',
+    lambda c: f'a bright photo of a {c}.',
+    lambda c: f'a photo of a clean {c}.',
+    lambda c: f'a photo of a dirty {c}.',
+    lambda c: f'a dark photo of the {c}.',
+    lambda c: f'a drawing of a {c}.',
+    lambda c: f'a photo of my {c}.',
+    lambda c: f'the plastic {c}.',
+    lambda c: f'a photo of the cool {c}.',
+    lambda c: f'a close-up photo of a {c}.',
+    lambda c: f'a black and white photo of the {c}.',
+    lambda c: f'a painting of the {c}.',
+    lambda c: f'a painting of a {c}.',
+    lambda c: f'a pixelated photo of the {c}.',
+    lambda c: f'a sculpture of the {c}.',
+    lambda c: f'a bright photo of the {c}.',
+    lambda c: f'a cropped photo of a {c}.',
+    lambda c: f'a plastic {c}.',
+    lambda c: f'a photo of the dirty {c}.',
+    lambda c: f'a jpeg corrupted photo of a {c}.',
+    lambda c: f'a blurry photo of the {c}.',
+    lambda c: f'a photo of the {c}.',
+    lambda c: f'a good photo of the {c}.',
+    lambda c: f'a rendering of the {c}.',
+    lambda c: f'a {c} in a video game.',
+    lambda c: f'a photo of one {c}.',
+    lambda c: f'a doodle of a {c}.',
+    lambda c: f'a close-up photo of the {c}.',
+    lambda c: f'a photo of a {c}.',
+    lambda c: f'the origami {c}.',
+    lambda c: f'the {c} in a video game.',
+    lambda c: f'a sketch of a {c}.',
+    lambda c: f'a doodle of the {c}.',
+    lambda c: f'a origami {c}.',
+    lambda c: f'a low resolution photo of a {c}.',
+    lambda c: f'the toy {c}.',
+    lambda c: f'a rendition of the {c}.',
+    lambda c: f'a photo of the clean {c}.',
+    lambda c: f'a photo of a large {c}.',
+    lambda c: f'a rendition of a {c}.',
+    lambda c: f'a photo of a nice {c}.',
+    lambda c: f'a photo of a weird {c}.',
+    lambda c: f'a blurry photo of a {c}.',
+    lambda c: f'a cartoon {c}.',
+    lambda c: f'art of a {c}.',
+    lambda c: f'a sketch of the {c}.',
+    lambda c: f'a embroidered {c}.',
+    lambda c: f'a pixelated photo of a {c}.',
+    lambda c: f'itap of the {c}.',
+    lambda c: f'a jpeg corrupted photo of the {c}.',
+    lambda c: f'a good photo of a {c}.',
+    lambda c: f'a plushie {c}.',
+    lambda c: f'a photo of the nice {c}.',
+    lambda c: f'a photo of the small {c}.',
+    lambda c: f'a photo of the weird {c}.',
+    lambda c: f'the cartoon {c}.',
+    lambda c: f'art of the {c}.',
+    lambda c: f'a drawing of the {c}.',
+    lambda c: f'a photo of the large {c}.',
+    lambda c: f'a black and white photo of a {c}.',
+    lambda c: f'the plushie {c}.',
+    lambda c: f'a dark photo of a {c}.',
+    lambda c: f'itap of a {c}.',
+    lambda c: f'graffiti of the {c}.',
+    lambda c: f'a toy {c}.',
+    lambda c: f'itap of my {c}.',
+    lambda c: f'a photo of a cool {c}.',
+    lambda c: f'a photo of a small {c}.',
+    lambda c: f'a tattoo of the {c}.',
+]
diff --git a/src/laion_clap/training/infer_demo.py b/src/laion_clap/training/infer_demo.py
index c346975..b39f3f4 100644
--- a/src/laion_clap/training/infer_demo.py
+++ b/src/laion_clap/training/infer_demo.py
@@ -1,8 +1,8 @@
 import torch
 import librosa
-from laion_clap import create_model
-from laion_clap.training.data import get_audio_features
-from laion_clap.training.data import int16_to_float32, float32_to_int16
+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
 
 tokenize = RobertaTokenizer.from_pretrained('roberta-base')
@@ -38,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():
     
@@ -63,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()
@@ -76,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/laion_clap/training/lp_main.py b/src/laion_clap/training/lp_main.py
index c2bd19c..1e88f73 100644
--- a/src/laion_clap/training/lp_main.py
+++ b/src/laion_clap/training/lp_main.py
@@ -24,16 +24,16 @@
 except ImportError:
     hvd = None
 
-from laion_clap import trace_model, create_model
-from laion_clap.training.data import get_data
-from laion_clap.training.params import parse_args
-from laion_clap.training.distributed import is_master, init_distributed_device, world_info_from_env
-from laion_clap.training.logger import setup_logging
-from laion_clap.training.scheduler import cosine_lr
-from laion_clap.training.lp_train import train_one_epoch, evaluate
-from laion_clap.utils import dataset_split, get_optimizer
-from laion_clap.utils import load_class_label
-from laion_clap.linear_probe import LinearProbe
+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 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/laion_clap/training/lp_train.py b/src/laion_clap/training/lp_train.py
index da5919c..d686336 100644
--- a/src/laion_clap/training/lp_train.py
+++ b/src/laion_clap/training/lp_train.py
@@ -14,8 +14,8 @@
 except ImportError:
     wandb = None
 
-from laion_clap import LPLoss, LPMetrics, lp_gather_features
-from laion_clap.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/laion_clap/training/main.py b/src/laion_clap/training/main.py
index 2d335e8..7c48b66 100644
--- a/src/laion_clap/training/main.py
+++ b/src/laion_clap/training/main.py
@@ -23,14 +23,14 @@
 except ImportError:
     hvd = None
 
-from laion_clap import trace_model, create_model
-from laion_clap.training.data import get_data
-from laion_clap.training.distributed import is_master, init_distributed_device, world_info_from_env
-from laion_clap.training.logger import setup_logging
-from laion_clap.training.params import parse_args
-from laion_clap.training.scheduler import cosine_lr
-from laion_clap.training.train import train_one_epoch, evaluate
-from laion_clap.utils import dataset_split, get_optimizer
+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 clap_module.utils import dataset_split, get_optimizer
 
 
 def maintain_ckpts(args, startidx, all_idx_len):
diff --git a/src/laion_clap/training/train.py b/src/laion_clap/training/train.py
index 07bfadd..0a28476 100644
--- a/src/laion_clap/training/train.py
+++ b/src/laion_clap/training/train.py
@@ -14,7 +14,7 @@
 except ImportError:
     wandb = None
 
-from laion_clap import ClipLoss, gather_features
+from clap_module import ClipLoss, gather_features
 from .distributed import is_master
 
 
@@ -597,7 +597,7 @@ def evaluate_clotho_audiocaps(
 
             # each item in the list has 5 texts
             if args.tmodel == "transformer":
-                from laion_clap 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/laion_clap/training/zero_shot.py b/src/laion_clap/training/zero_shot.py
new file mode 100644
index 0000000..04472c1
--- /dev/null
+++ b/src/laion_clap/training/zero_shot.py
@@ -0,0 +1,90 @@
+# NOTE: This script is currently not supported for CLAP.
+import logging
+from contextlib import suppress
+
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+from clap_module import tokenize
+from .imagenet_zeroshot_data import imagenet_classnames, openai_imagenet_template
+
+
+def zero_shot_classifier(model, classnames, templates, args):
+    with torch.no_grad():
+        zeroshot_weights = []
+        for classname in tqdm(classnames):
+            texts = [template(classname) for template in templates]  # format with class
+            texts = tokenize(texts).to(args.device)  # tokenize
+            if args.distributed and not args.horovod:
+                class_embeddings = model.module.encode_text(texts)
+            else:
+                class_embeddings = model.encode_text(texts)
+            class_embedding = F.normalize(class_embeddings, dim=-1).mean(dim=0)
+            class_embedding /= class_embedding.norm()
+            zeroshot_weights.append(class_embedding)
+        zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(args.device)
+    return zeroshot_weights
+
+
+def accuracy(output, target, topk=(1,)):
+    pred = output.topk(max(topk), 1, True, True)[1].t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+    return [float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy()) for k in topk]
+
+
+def run(model, classifier, dataloader, args):
+    autocast = torch.cuda.amp.autocast if args.precision == 'amp' else suppress
+    with torch.no_grad():
+        top1, top5, n = 0., 0., 0.
+        for images, target in tqdm(dataloader, unit_scale=args.batch_size):
+            images = images.to(args.device)
+            target = target.to(args.device)
+
+            with autocast():
+                # predict
+                if args.distributed and not args.horovod:
+                    image_features = model.module.encode_image(images)
+                else:
+                    image_features = model.encode_image(images)
+                image_features = F.normalize(image_features, dim=-1)
+                logits = 100. * image_features @ classifier
+
+            # measure accuracy
+            acc1, acc5 = accuracy(logits, target, topk=(1, 5))
+            top1 += acc1
+            top5 += acc5
+            n += images.size(0)
+
+    top1 = (top1 / n)
+    top5 = (top5 / n)
+    return top1, top5
+
+
+def zero_shot_eval(model, data, epoch, args):
+    if 'imagenet-val' not in data and 'imagenet-v2' not in data:
+        return {}
+    if args.zeroshot_frequency == 0:
+        return {}
+    if (epoch % args.zeroshot_frequency) != 0 and epoch != args.epochs:
+        return {}
+
+    logging.info('Starting zero-shot imagenet.')
+
+    logging.info('Building zero-shot classifier')
+    classifier = zero_shot_classifier(model, imagenet_classnames, openai_imagenet_template, args)
+
+    logging.info('Using classifier')
+    results = {}
+    if 'imagenet-val' in data:
+        top1, top5 = run(model, classifier, data['imagenet-val'].dataloader, args)
+        results['imagenet-zeroshot-val-top1'] = top1
+        results['imagenet-zeroshot-val-top5'] = top5
+    if 'imagenet-v2' in data:
+        top1, top5 = run(model, classifier, data['imagenet-v2'].dataloader, args)
+        results['imagenetv2-zeroshot-val-top1'] = top1
+        results['imagenetv2-zeroshot-val-top5'] = top5
+
+    logging.info('Finished zero-shot imagenet.')
+
+    return results


Test patch:
diff --git a/src/laion_clap/unit_test.py b/src/laion_clap/unit_test.py
index c1de796..491db7d 100644
--- a/src/laion_clap/unit_test.py
+++ b/src/laion_clap/unit_test.py
@@ -30,7 +30,13 @@
 print(audio_embed.shape)
 
 # Get text embedings from texts:
-text_data = ["I love the contrastive learning", "I love the pretrain model"]
+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)
\ No newline at end of file
+print(text_embed.shape)
+
+
+
+
+
+