update

Dockerfile

@@ -36,7 +36,8 @@ RUN pip install streamlit --timeout 500
 RUN pip install ffmpeg-python --timeout 1000
 RUN pip install toml
 RUN pip install librosa
-
+RUN pip install pandas
+RUN pip install pyannote-audio
 # RUN pip uninstall ffmpeg --yes
 # RUN pip uninstall ffmpeg-python --yes
 

app.py

@@ -7,6 +7,8 @@ import whisper
 from sklearn.cluster import AgglomerativeClustering
 import torch
 import librosa
+from mix_sae import *
+
 
 UPLOAD_FOLDER = "./uploads"
 

app_test.py

@@ -0,0 +1,191 @@
+import streamlit as st
+import numpy as np
+import os
+import whisper
+from sklearn.cluster import AgglomerativeClustering
+import torch
+import librosa
+from torch.utils.data import DataLoader
+from mix_sae import MoESparseAutoencodersCL
+from load_dataset import AutoEncoderDataset
+import argparse
+
+
+UPLOAD_FOLDER = "./uploads"
+parser = argparse.ArgumentParser(description='Deep Clustering Network')
+parser.add_argument('--input_dim', type=int, default=384,
+                    help='input dimension')
+# Model parameters
+parser.add_argument('--lr', type=float, default=1e-3,
+                    help='learning rate (default: 1e-4)')
+parser.add_argument('--wd', type=float, default=1e-4,
+                    help='weight decay (default: 5e-4)')
+parser.add_argument('--batch-size', type=int, default=16,
+                    help='input batch size for training')
+parser.add_argument('--lamda', type=float, default=1,
+                    help='coefficient of the reconstruction loss')
+parser.add_argument('--beta', type=float, default=1,
+                    help=('coefficient of the regularization term on '
+                            'clustering'))
+parser.add_argument('--hidden-dims', default=[256, 128, 64, 32],
+                    help='learning rate (default: 1e-4)')
+parser.add_argument('--latent_dim', type=int, default=2,
+                    help='latent space dimension')
+parser.add_argument('--n-clusters', type=int, default=2,
+                    help='number of clusters in the latent space')
+parser.add_argument('--input-dim', type=int, default=384,
+                    help='input dimension')
+parser.add_argument('--n-classes', type=int, default=2,
+                    help='output dimension')
+parser.add_argument('--pretrain_epochs', type=int, default=30,
+                    help='pretraining step epochs')
+parser.add_argument('--pretrain_epochs_main', type=int, default=30,
+                    help='pretraining step epochs')
+parser.add_argument('--pretrain', type=bool, default=True,
+                    help='whether use pre-training')
+parser.add_argument('--main_train_epochs', type=int, default=5,
+                    help='main_train epochs')
+parser.add_argument('--rho', type=float, default=0.2,
+                    help='whether use pre-training')
+parser.add_argument('--sparsity_param', type=float, default=0.1,
+                    help='sparsity constract param')
+parser.add_argument('--cl_loss_param', type=float, default=0.05,
+                    help='clasification loss param')
+args = parser.parse_args()
+
+
+def allowed_file(filename):
+    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ['wav']
+
+
+def process_wav(audio_file, speaker_number, model_type, run_device = 'cpu', sr = 16000):
+    embedding_dims = {"tiny": 384, 'small': 768, 'base': 512, 'medium':1024}
+    #---- get results from whisper model
+    whisper_model = whisper.load_model(model_type, run_device)
+    wp_results = whisper_model.transcribe(audio_file)
+    for ide in range(len(wp_results['segments'])):
+        del wp_results['segments'][ide]['seek']
+        del wp_results['segments'][ide]['tokens']
+        del wp_results['segments'][ide]['compression_ratio']
+        del wp_results['segments'][ide]['temperature']
+        del wp_results['segments'][ide]['avg_logprob']
+        del wp_results['segments'][ide]['no_speech_prob']
+
+    #---- solve each segment
+    segments = wp_results["segments"]
+
+    # >= 2 sentences
+    if len(segments) > 1:
+        embeddings = np.zeros(shape=(len(segments), embedding_dims[model_type]))
+
+        for i, segment in enumerate(segments):
+            start = int(segment["start"] * sr)
+            end = int(segment["end"] * sr)
+
+            # Extract a segment
+            audio = audio_file[start: end]
+            mel = whisper.log_mel_spectrogram(audio).to(whisper_model.device)
+            
+            #--- this code to create the correct shape of mel spectrogram
+            while True:
+                nF, nT = np.shape(mel)
+                if nT > 3000:
+                    mel = mel[:,0:3000]
+                    break
+                else:
+                    mel = torch.cat((mel, mel), -1)
+            mel = torch.unsqueeze(mel, 0)
+            wp_emb = whisper_model.embed_audio(mel)
+            #print(np.shape(wp_emb))
+
+            emb_1d  = np.mean(wp_emb.cpu().detach().numpy(), axis=0)
+            emb_1d  = np.mean(emb_1d, axis=0)
+            #print(np.shape(emb_1d))
+            #exit()
+            embeddings[i] = emb_1d
+        embeddings= np.array(embeddings, dtype="f")
+        train_loader = AutoEncoderDataset(embeddings)
+        train_loader = DataLoader(train_loader, batch_size = args.batch_size, shuffle = False)
+
+    
+        moe_cl = MoESparseAutoencodersCL(args=args)
+        mixture_moe_cl, full_latent_X = moe_cl.pretraining(train_loader)
+        pre_label = moe_cl.psedo_label 
+        mixture_moe_cl = moe_cl.main_training(train_loader)
+        moe_cl_pred = moe_cl.get_final_cluster(train_loader)
+        #--- clustering spk emb
+    #     clustering = AgglomerativeClustering(speaker_number, compute_distances=True).fit(embeddings)
+    #     labels = clustering.labels_
+
+        for i in range(len(segments)):
+            wp_results['segments'][i]["speaker"] = 'SPEAKER ' + str(pre_label[i] + 1)
+
+    # only one sentence
+    else:
+         wp_results['segments'][0]["speaker"] = 'SPEAKER 1'
+
+    return wp_results 
+
+
+def main():
+
+    title_style = """
+    <style>
+    .title {
+        text-align: center;
+        font-size: 40px;
+    }
+    </style>
+    """
+    st.markdown(
+    title_style,
+    unsafe_allow_html=True
+    )
+    title  = """
+    <h1 class = "title" >Telephone Calls Speaker Diarization</h1>
+    </div>
+    """
+    st.markdown(title,
+                unsafe_allow_html=True)
+    # st.title("Speaker Diarization")
+
+
+    # Get user inputs
+    file = st.file_uploader("Upload a WAV file:", type=["wav"])
+    num_speakers = st.number_input("Number of speakers:", min_value=2, max_value=2)
+
+    model_list = ['tiny', 'small', 'base', 'medium']
+    model_type = st.selectbox("Select model type: ", model_list)
+
+    # Display the result
+    st.write("Your uploaded wav file: ")
+    st.audio(file, format = 'audio/wav')
+    if st.button("Submit"):
+        if file is not None:
+            
+            # Read audio file using pydub
+            audio_file, _ = librosa.load(file, sr=16000)
+
+            # Process the uploaded file using the AI model
+            wp_results = process_wav(audio_file, num_speakers, model_type)
+
+            # Write result:
+            st.write("Segments:" )
+            for seg in wp_results['segments']:
+                seg['start'] = np.round(seg['start'], 1)
+                seg['end'] = np.round(seg['end'], 1)
+                st.write(seg)
+            st.write("Language: ", wp_results['language'])
+            st.write("Full text:")
+            st.write(wp_results['text'])
+        else:
+            print("Error")
+    st.write("\n\n---\n\n")  
+    st.write("Built with Docker and Streamlit")
+    st.link_button("Paper link: https://arxiv.org/abs/2407.01963", "https://arxiv.org/abs/2407.01963")
+    return 
+
+
+
+if __name__ == "__main__":
+    main()

create_DER.py

@@ -0,0 +1,232 @@
+import numpy as np
+import pandas as pd
+import argparse
+import os
+import simpleder
+from pyannote.metrics.diarization import DiarizationErrorRate
+from pyannote.core import Segment, Annotation
+
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--label_dir", type = str, default="./datasets/spanish/human_label", help= "rttm label dir")
+
+opt = parser.parse_args()
+LABEL_DIR = opt.label_dir
+column = ['speaker','file_name','number' ,'start', 'duration', 'na1', 'na2', 'label', 'na3', 'na4']
+
+
+def createDER(label_path, sample_dir, prediction, window_length = 0.5, overlap = 0.0):
+    """
+    Extract series from label and prediction for calculating DER
+    """
+    
+    # df = pd.read_csv(label_path, delimiter=' ', header=None, usecols=column, names=column)
+    # ref = []
+    
+    # prev_end = 0
+    # # Assign label
+    # for row in df.iterrows():
+    #     row_item = row[1]
+    #     start = np.round(row_item['start'], 2)
+    #     end = np.round(row_item['start'] + row_item['duration'], 2)
+    #     # Avoid overlap
+    #     if start < prev_end:
+    #         start = prev_end
+    #     # Avoid error label
+    #     if start > end:
+    #         continue
+    #     ref.append((row_item['label'], start, end))
+    #     prev_end = end
+
+    
+    df = pd.read_csv(label_path, delimiter=' ', header=None, usecols=column, names=column)
+    refer = Annotation(uri='label')
+
+    # Assign label
+    prev_end = 0
+    for row in df.iterrows():
+        row_item = row[1]
+        start = np.round(row_item['start'], 2)
+        end = np.round(row_item['start'] + row_item['duration'], 2)
+        # Avoid overlap
+        if start < prev_end:
+            start = prev_end
+        # Avoid error label
+        if start > end:
+            continue
+        refer[Segment(start, end)] = row_item['label']
+        prev_end = end
+    print("******EXTRACT LABEL DONE***********")
+
+    # assert len(os.listdir(sample_dir)) == len(prediction)
+    segment_list = sorted(os.listdir(sample_dir), key= lambda x: float(x.split("_")[-2]))
+
+    # Create index mapping to store start-end index of consecutive segments
+    index_mapping = {}
+    start_index = 0
+    current_value = prediction[0]
+    
+    for i in range(1, len(prediction)):
+        if prediction[i] != current_value:
+            index_mapping[(start_index, i - 1)] = current_value
+            start_index = i
+            current_value = prediction[i]
+    
+    # Handle the last consecutive sequence
+    index_mapping[(start_index, len(prediction) - 1)] = current_value
+    
+
+    # Assign label to consecutive segments
+    hyp = []
+    for key, value in index_mapping.items():
+        start_index = key[0]
+        end_index = key[1]
+        speaker_label = "spk0{}".format(value)
+        if overlap != 0:
+            start_time = np.round(overlap * start_index, 2)
+            if start_index == end_index:
+                end_time = np.round(start_time + window_length, 2)
+            else:
+                end_time = np.round(overlap * end_index + window_length, 2)
+        # Non-overlap
+        else:
+            start_time = np.round(window_length * start_index, 2)
+            if start_index == end_index:
+                end_time = np.round(start_time + window_length, 2)
+            else:
+                end_time = np.round((end_index + 1) * window_length, 2)
+
+        hyp.append((speaker_label, start_time, end_time))
+
+
+    hypo = Annotation(uri='hypo')
+    for item in hyp:
+        hypo[Segment(item[1], item[2])] = item[0]
+
+    print("******EXTRACT HYP DONE***********")
+    
+    return refer, hypo
+
+
+def create_DER_pyannote(label_path, pyannote_label_path):
+    # df = pd.read_csv(label_path, delimiter=' ', header=None, usecols=column, names=column)
+
+    # ref = []
+    # prev_end = 0
+    # # Assign label
+    # for row in df.iterrows():
+    #     row_item = row[1]
+    #     start = np.round(row_item['start'], 2)
+    #     end = np.round(row_item['start'] + row_item['duration'], 2)
+    #     # Avoid overlap
+    #     if start < prev_end:
+    #         start = prev_end
+    #     # Avoid error label
+    #     if start > end:
+    #         continue
+    #     ref.append((row_item['label'], start, end))
+    #     prev_end = end
+    df = pd.read_csv(label_path, delimiter=' ', header=None, usecols=column, names=column)
+    refer = Annotation(uri='label')
+
+    # Assign label
+    for row in df.iterrows():
+        row_item = row[1]
+        start = np.round(row_item['start'], 2)
+        end = np.round(row_item['start'] + row_item['duration'], 2)
+        # # Avoid overlap
+        # if start < prev_end:
+        #     start = prev_end
+        # # Avoid error label
+        # if start > end:
+        #     continue
+        refer[Segment(start, end)] = row_item['label']
+        # ref.append((row_item['label'], start, end))
+        # prev_end = end
+
+    print("******EXTRACT LABEL DONE*****c******")
+
+    df = pd.read_csv(pyannote_label_path, delimiter=' ', header=None, usecols=column, names=column)
+    print(df)
+    pyannote_ref = []
+    prev_end = 0
+    # Assign label
+    for row in df.iterrows():
+        row_item = row[1]
+        start = np.round(row_item['start'], 2)
+        end = np.round(row_item['start'] + row_item['duration'], 2)
+        # Avoid overlap
+        if start < prev_end:
+            start = prev_end
+        # Avoid error label
+        if start > end:
+            continue
+        pyannote_ref.append((row_item['label'], start, end))
+        prev_end = end
+    print("******EXTRACT PYANNOTE LABEL DONE***********")
+    return refer, pyannote_ref
+
+
+def create_pyannote_timeline(label_path, pyannote_label_path):
+    df = pd.read_csv(label_path, delimiter=' ', header=None, usecols=column, names=column)
+    refer = Annotation(uri='label')
+    # ref = []
+    # prev_end = 0
+    # Assign label
+    for row in df.iterrows():
+        row_item = row[1]
+        start = np.round(row_item['start'], 2)
+        end = np.round(row_item['start'] + row_item['duration'], 2)
+        # # Avoid overlap
+        # if start < prev_end:
+        #     start = prev_end
+        # # Avoid error label
+        # if start > end:
+        #     continue
+        refer[Segment(start, end)] = row_item['label']
+        # ref.append((row_item['label'], start, end))
+        # prev_end = end
+
+    print("******EXTRACT LABEL DONE***********")
+
+    df = pd.read_csv(pyannote_label_path, delimiter=' ', header=None, usecols=column, names=column)
+    py_refer = Annotation(uri='py_label')
+    ref = []
+    # prev_end = 0
+    # Assign label
+    for row in df.iterrows():
+        row_item = row[1]
+        start = np.round(row_item['start'], 2)
+        end = np.round(row_item['start'] + row_item['duration'], 2)
+        # # Avoid overlap
+        # if start < prev_end:
+        #     start = prev_end
+        # # Avoid error label
+        # if start > end:
+        #     continue
+        py_refer[Segment(start, end)] = row_item['label']
+        # ref.append((row_item['label'], start, end))
+        # prev_end = end
+
+    print("******EXTRACT PY LABEL DONE***********")
+
+    return refer, py_refer
+
+
+
+if __name__ == "__main__":
+    label_dir = "datasets/spanish/human_label"
+    py_label_dir = "datasets/spanish/label"
+    label_list = sorted(os.listdir(label_dir))[:11]
+    py_label_list = sorted(os.listdir(py_label_dir))[:11]
+    with open("./compare_py_label_PY.txt", "w") as file:
+        for label, py_label in zip(label_list, py_label_list):
+            label_path = label_dir + "/" + label
+            py_label_path = py_label_dir + '/' + py_label
+            ref, py_ref = create_pyannote_timeline(label_path=label_path, pyannote_label_path=py_label_path)
+            der = DiarizationErrorRate(collar=0.0, skip_overlap=False)
+            error = der(ref, py_ref)
+            file.write(str(label) + "PYANNOTE err:" + str(error) + "\n")
+    file.close()
+

load_dataset.py

@@ -0,0 +1,114 @@
+import torch
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+import numpy as np
+import os
+
+
+class CustomDataset(torch.utils.data.Dataset):
+    def __init__(self, sample_dir, embed_dim = 384, train = False, time = 2):
+        self.sample_dir = sample_dir
+        self.n_segments = len(os.listdir(self.sample_dir))
+        self.data = np.zeros((self.n_segments, embed_dim))
+        
+        # Sorted segment based on start_time
+        self.sorted_segments = sorted(os.listdir(sample_dir), key= lambda x: float(x.split("_")[-2]))
+
+        # Assign segments
+        for idx, segment_npy in enumerate(self.sorted_segments):
+            segment_path = self.sample_dir + "/" + segment_npy
+            segment_embed = np.load(segment_path)
+            self.data[idx] = segment_embed
+            
+        if train:
+            for time in range(time):
+                self.data = np.concatenate((self.data, self.data), axis = 0)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        sample = torch.from_numpy(self.data[idx]).float()
+        return sample
+    
+
+class AutoEncoderDataset(torch.utils.data.Dataset):
+    """
+    Create dataset from predefined tensor for each autoencoder in MOE
+    """
+    def __init__(self, data):
+        self.data = data
+    def __len__(self):
+        return len(self.data)
+    def __getitem__(self, idx):
+        return self.data[idx]
+    
+    
+import argparse
+parser = argparse.ArgumentParser(description='Deep Clustering Network')
+
+# Dataset parameters
+parser.add_argument('--dir', default='./datasets/spanish/',
+                    help='dataset directory')
+parser.add_argument('--input_dim', type=int, default=384,
+                    help='input dimension')
+parser.add_argument('--n-classes', type=int, default=2,
+                    help='output dimension')
+
+# Training parameters
+parser.add_argument('--lr', type=float, default=1e-3,
+                    help='learning rate (default: 1e-4)')
+parser.add_argument('--wd', type=float, default=1e-4,
+                    help='weight decay (default: 5e-4)')
+parser.add_argument('--batch-size', type=int, default=16,
+                    help='input batch size for training')
+parser.add_argument('--epoch', type=int, default=50,
+                    help='number of epochs to train')
+parser.add_argument('--pre-epoch', type=int, default=100, 
+                    help='number of pre-train epochs')
+parser.add_argument('--pretrain', type=bool, default=True,
+                    help='whether use pre-training')
+
+# Model parameters
+parser.add_argument('--lamda', type=float, default=1,
+                    help='coefficient of the reconstruction loss')
+parser.add_argument('--beta', type=float, default=1, 
+                    help=('coefficient of the regularization term on '
+                            'clustering'))
+parser.add_argument('--hidden-dims', default=[256, 128, 64, 32, 16],
+                    help='learning rate (default: 1e-4)')
+parser.add_argument('--latent_dim', type=int, default=2,
+                    help='latent space dimension')
+parser.add_argument('--n-clusters', type=int, default=2,
+                    help='number of clusters in the latent space')
+
+parser.add_argument('--n_1Dconv', type=int, default=4,
+                    help='n_1dconv')
+parser.add_argument('--kernel_size', default=[7, 5, 3, 3],
+                    help='kernel_size')
+parser.add_argument('--stride', type = int, default=1,
+                    help='stride')
+parser.add_argument('--num_blocks', type = int, default=4,
+                    help='num_blocks')
+parser.add_argument('--channels', type = int, default=[128, 64, 32, 16],
+                    help='channels')
+
+# Utility parameters
+parser.add_argument('--n-jobs', type=int, default=1,
+                    help='number of jobs to run in parallel')
+parser.add_argument('--log-interval', type=int, default=20,
+                    help=('how many batches to wait before logging the '
+                            'training status'))
+parser.add_argument("--window_length", type = float, default= 0.4, help="window length")
+parser.add_argument("--overlap", type = float, default= 0, help="overlap")
+
+
+args = parser.parse_args()
+
+if __name__ == "__main__":
+    # Example usage:
+    sample_dir = "datasets/spanish/segments/0096_[cut_193sec].wav"
+    dataset = CustomDataset(sample_dir=sample_dir,train= False)
+    # dataset = CustomDataset(sample_dir=sample_dir,train= False)
+
+    

mix_sae.py

@@ -0,0 +1,672 @@
+import torch.nn as nn
+from collections import OrderedDict
+import torch
+import argparse
+import torch.nn.init as init
+import numpy as np
+from sklearn.cluster import KMeans, SpectralClustering
+from load_dataset import AutoEncoderDataset
+from torch.utils.data import DataLoader
+from load_dataset import *
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+#----------------------VERSION 3: SPARSE AUTOENDCODER KL PENALTY AND ENTROPY LOSS--------------------
+import random
+
+random.seed(10)
+
+parser = argparse.ArgumentParser(description='Deep Clustering Network')
+parser.add_argument('--input_dim', type=int, default=384,
+                    help='input dimension')
+# Model parameters
+parser.add_argument('--lr', type=float, default=1e-3,
+                    help='learning rate (default: 1e-4)')
+parser.add_argument('--wd', type=float, default=1e-4,
+                    help='weight decay (default: 5e-4)')
+parser.add_argument('--batch-size', type=int, default=16,
+                    help='input batch size for training')
+parser.add_argument('--lamda', type=float, default=1,
+                    help='coefficient of the reconstruction loss')
+parser.add_argument('--beta', type=float, default=1,
+                    help=('coefficient of the regularization term on '
+                            'clustering'))
+parser.add_argument('--hidden-dims', default=[256, 128, 64, 32],
+                    help='learning rate (default: 1e-4)')
+parser.add_argument('--latent_dim', type=int, default=2,
+                    help='latent space dimension')
+parser.add_argument('--n-clusters', type=int, default=2,
+                    help='number of clusters in the latent space')
+parser.add_argument('--input-dim', type=int, default=384,
+                    help='input dimension')
+parser.add_argument('--n-classes', type=int, default=2,
+                    help='output dimension')
+parser.add_argument('--pretrain_epochs', type=int, default=80,
+                    help='pretraining step epochs')
+parser.add_argument('--pretrain_epochs_main', type=int, default=80,
+                    help='pretraining step epochs')
+parser.add_argument('--pretrain', type=bool, default=True,
+                    help='whether use pre-training')
+parser.add_argument('--main_train_epochs', type=int, default=80,
+                    help='main_train epochs')
+parser.add_argument('--rho', type=float, default=0.2,
+                    help='whether use pre-training')
+parser.add_argument('--sparsity_param', type=float, default=0.1,
+                    help='sparsity constract param')
+parser.add_argument('--cl_loss_param', type=float, default=0.05,
+                    help='clasification loss param')
+args = parser.parse_args()
+
+
+
+
+class AutoEncoder(nn.Module):
+
+    def __init__(self, args):
+        super(AutoEncoder, self).__init__()
+        self.args = args
+        self.input_dim = args.input_dim
+        self.output_dim = self.input_dim
+        self.hidden_dims = args.hidden_dims
+        self.hidden_dims.append(args.latent_dim)
+        self.dims_list = (args.hidden_dims +
+                          args.hidden_dims[:-1][::-1])  # mirrored structure
+        self.n_layers = len(self.dims_list)
+        self.latent_dim = args.latent_dim
+        self.n_clusters = args.n_clusters
+        self.RHO = args.rho
+
+        # Validation check
+        assert self.n_layers % 2 > 0
+        assert self.dims_list[self.n_layers // 2] == self.latent_dim
+
+        # Encoder Network
+        layers = OrderedDict()
+        for idx, hidden_dim in enumerate(self.hidden_dims):
+            if idx == 0:
+                layers.update(
+                    {
+                        'linear0': nn.Linear(self.input_dim, hidden_dim),
+                        # 'linear0': CustomDense(self.input_dim, hidden_dim),
+                        # 'activation0': nn.LeakyReLU()
+                        # 'activation0': nn.ReLU()
+                    }
+                )
+            else:
+                layers.update(
+                    {
+                        'linear{}'.format(idx): nn.Linear(
+                            self.hidden_dims[idx-1], hidden_dim),
+                        # 'linear{}'.format(idx): CustomDense(self.hidden_dims[idx-1], hidden_dim),
+                        # 'activation{}'.format(idx): nn.LeakyReLU(),
+                        # 'activation{}'.format(idx): nn.ELU(),
+                        'activation{}'.format(idx): nn.LeakyReLU(),
+                        # 'dropout{}'.format(idx): nn.Dropout(0.5), 
+                        'bn{}'.format(idx): nn.BatchNorm1d(
+                            self.hidden_dims[idx]),
+                        
+
+                        # 'bn{}'.format(idx): nn.BatchNorm1d(
+                        #     self.hidden_dims[idx])
+                    }
+                )
+        self.encoder = nn.Sequential(layers)
+
+        # Decoder Network
+        layers = OrderedDict()
+        tmp_hidden_dims = self.hidden_dims[::-1]
+        for idx, hidden_dim in enumerate(tmp_hidden_dims):
+            if idx == len(tmp_hidden_dims) - 1:
+                layers.update(
+                    {
+                        'linear{}'.format(idx): nn.Linear(
+                            hidden_dim, self.output_dim),
+                        # 'activation{}'.format(idx):nn.ReLU()
+                        # 'activation{}'.format(idx): nn.LeakyReLU(),
+                        # 'activation{}'.format(idx): nn.ELU(),
+                        #  'linear{}'.format(idx): CustomDense(hidden_dim, self.output_dim),
+                    }
+                )
+            else:
+                layers.update(
+                    {
+                        'linear{}'.format(idx): nn.Linear(
+                            hidden_dim, tmp_hidden_dims[idx+1]),
+                        # 'linear{}'.format(idx): CustomDense(
+                        #     hidden_dim, tmp_hidden_dims[idx+1]),
+                        # 'activation{}'.format(idx): nn.ELU(),
+                        'activation{}'.format(idx): nn.LeakyReLU(),
+                        # 'dropout{}'.format(idx): nn.Dropout(0.5),
+                        'bn{}'.format(idx): nn.BatchNorm1d(
+                            tmp_hidden_dims[idx+1]),
+                        # 'activation{}'.format(idx): nn.ELU(),
+                        
+                        # 'bn{}'.format(idx): nn.BatchNorm1d(
+                        #     tmp_hidden_dims[idx+1])
+                    }
+                )
+        self.decoder = nn.Sequential(layers)
+        # Apply Xavier weight initialization to all linear layers
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                init.xavier_normal_(m.weight)
+                init.constant_(m.bias, 0)  # Initialize biases to 0
+    def __repr__(self):
+        repr_str = '[Structure]: {}-'.format(self.input_dim)
+        for idx, dim in enumerate(self.dims_list):
+            repr_str += '{}-'.format(dim)
+        repr_str += str(self.output_dim) + '\n'
+        repr_str += '[n_layers]: {}'.format(self.n_layers) + '\n'
+        repr_str += '[n_clusters]: {}'.format(self.n_clusters) + '\n'
+        repr_str += '[input_dims]: {}'.format(self.input_dim)
+        return repr_str
+
+    def __str__(self):
+        return self.__repr__()
+
+    def forward(self, X, latent=False):
+        output = self.encoder(X)
+        if latent:
+            return output
+        return self.decoder(output)
+
+
+class VAE(nn.Module):
+    def __init__(self, args):
+        super(VAE, self).__init__()
+        self.args = args
+        self.input_dim = args.input_dim
+        self.output_dim = self.input_dim
+        self.hidden_dims = args.hidden_dims
+        self.latent_dim = args.latent_dim
+        self.n_clusters = args.n_clusters
+
+        # Encoder Network
+        layers = OrderedDict()
+        for idx, hidden_dim in enumerate(self.hidden_dims):
+            if idx == 0:
+                layers.update(
+                    {
+                        'linear0': nn.Linear(self.input_dim, hidden_dim),
+                      
+                    }
+                )
+            else:
+                layers.update(
+                    {
+                        'linear{}'.format(idx): nn.Linear(
+                            self.hidden_dims[idx-1], hidden_dim),
+                       
+                        'activation{}'.format(idx): nn.ReLU(),
+
+                        'bn{}'.format(idx): nn.BatchNorm1d(
+                            self.hidden_dims[idx])
+                    }
+                )
+        self.encoder = nn.Sequential(layers)
+
+        # Decoder Network
+        layers = OrderedDict()
+        tmp_hidden_dims = self.hidden_dims[::-1]
+        for idx, hidden_dim in enumerate(tmp_hidden_dims):
+            if idx == len(tmp_hidden_dims) - 1:
+                layers.update(
+                    {
+                        'linear{}'.format(idx): nn.Linear(
+                            hidden_dim, self.output_dim),
+                    }
+                )
+            else:
+                layers.update(
+                    {
+                        'linear{}'.format(idx): nn.Linear(
+                            hidden_dim, tmp_hidden_dims[idx+1]),
+                        'activation{}'.format(idx): nn.ReLU(),
+                        'bn{}'.format(idx): nn.BatchNorm1d(
+                            tmp_hidden_dims[idx+1])
+                    }
+                )
+        self.decoder = nn.Sequential(layers)
+        self.fc_mu = nn.Linear(self.hidden_dims[-1], self.latent_dim)
+        self.fc_var = nn.Linear(self.hidden_dims[-1], self.latent_dim)
+
+        self.decode_input_linear = nn.Linear(self.latent_dim, self.hidden_dims[-1])
+        # Apply Xavier weight initialization to all linear layers
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                init.xavier_normal_(m.weight)
+                init.constant_(m.bias, 0)  # Initialize biases to 0
+
+    def encode(self, x):
+        x = self.encoder(x)
+        mu = self.fc_mu(x)
+        log_var = self.fc_var(x)
+
+        return [mu, log_var]
+    
+    def decode(self, x):
+        x = self.decode_input_linear(x)
+        x = self.decoder(x)
+        return x
+
+    def reparameterize(self, mu, logvar):
+        """
+        Reparameterization trick to sample from N(mu, var) from
+        N(0,1).
+        """
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+    
+
+    def loss_function(self, x_hat, x, mu, log_var, kld_weight = 1):
+        """
+        Computes the VAE loss function.
+        KL(N(\mu, \sigma), N(0, 1)) = \log \frac{1}{\sigma} + \frac{\sigma^2 + \mu^2}{2} - \frac{1}{2}
+        """
+
+        rec_loss = torch.nn.functional.mse_loss(x_hat, x)
+
+        kld_loss = torch.mean(-0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim = 1), dim = 0)
+
+        loss = rec_loss + kld_weight * kld_loss
+
+        return [loss, rec_loss.detach(), -kld_loss.detach()]
+
+    def forward(self, x):
+        """
+        Forward VAE
+        Return: [output, input, mu, var]
+        """
+        # Encoder
+        mu, log_var = self.encode(x) 
+        # Sample
+        z = self.reparameterize(mu, log_var)
+        # Decoder
+        output = self.decode(z)
+
+        return [output, x, mu, log_var]
+    
+
+class ClusterNet(nn.Module):
+
+    def __init__(self, input_dim, hidden_dims = [128], n_clusters=2):
+        """ClusterNet("""
+        super(ClusterNet, self).__init__()
+        layers = []
+        for i in range(len(hidden_dims)):
+            if i == 0:
+                layers.append(nn.Linear(input_dim, hidden_dims[i]))
+                layers.append(nn.LeakyReLU())
+                # layers.append(nn.Dropout(0.5))
+                # layers.append(nn.BatchNorm1d(hidden_dims[i])),
+            else:
+                layers.append(nn.Linear(hidden_dims[i-1], hidden_dims[i])),
+                layers.append(nn.LeakyReLU())
+                # layers.append(nn.Dropout(0.5))
+                # layers.append(nn.BatchNorm1d(hidden_dims[i])),
+        # Last layer
+        layers.append(nn.Sequential(
+            nn.Flatten(), 
+            nn.Linear(hidden_dims[-1], n_clusters),
+            nn.Softmax(dim = 1),
+        ))
+
+        self.layers = nn.Sequential(*layers)
+        # Apply Xavier weight initialization to all linear layers
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                init.xavier_normal_(m.weight)
+                init.constant_(m.bias, 0)  # Initialize biases to 0
+    def forward(self, x):
+        """Extract the feature vectors."""
+        features = x
+        for layer in self.layers:
+            features = layer(features)
+        return features
+    
+
+class MoESparseAutoencodersCL(nn.Module):
+    """
+    Mixture of Expert DNN-Autoencoder
+    """
+    def __init__(self, args):
+        super(MoESparseAutoencodersCL, self).__init__()
+        self.args = args
+        self.input_dim = args.input_dim
+        self.output_dim = self.input_dim
+        self.hidden_dims = args.hidden_dims
+        self.latent_dim = args.latent_dim
+        self.n_clusters = args.n_clusters
+        self.pretrain_epochs = args.pretrain_epochs
+        self.pretrain_epochs_main = args.pretrain_epochs_main
+        self.main_train_epochs = args.main_train_epochs
+        self.device = "cpu"
+        # Define main autoencoder at pretraining
+        self.main_autoencoder = AutoEncoder(args=args)
+        self.RHO = args.rho
+        self.BETA = args.sparsity_param
+        self.psedo_label = None
+        # Clustering algorithm for pre-training
+        self.cluster_algo = None
+        self.cl_loss_param = args.cl_loss_param
+
+        # Define autoencoder expert in mixture
+        self.moe = {}
+        for i in range(self.n_clusters):
+            self.moe[i] = AutoEncoder(args)
+        # Add cluster net (gating network) to moe 
+        self.moe['cluster_net'] = ClusterNet(input_dim= self.input_dim, n_clusters=self.n_clusters)
+    
+    
+    def kl_divergence(self, rho, rho_hat):
+        rho_hat = torch.mean(F.sigmoid(rho_hat), 1) # sigmoid because we need the probability distributions
+        rho = torch.tensor([rho] * len(rho_hat)).to(self.device)
+        return torch.sum(rho * torch.log(rho/rho_hat) + (1 - rho) * torch.log((1 - rho)/(1 - rho_hat)))
+    
+    # define the sparse loss function
+    def sparse_loss(self, rho, X, model):
+        values = X
+        loss = 0
+        model_children = list(model.children())
+        for i in range(len(model_children)):
+            values = model_children[i](values)
+            loss += self.kl_divergence(rho, values)
+        return loss / X.shape[0]
+    
+
+
+    def batchwise_entropy_loss(self, cluster_outputs):
+        """
+        Calculate batch wise entropy loss
+        """
+        X = torch.mean(cluster_outputs, axis = 0)
+        return torch.special.entr(X).sum()
+        
+
+
+    def loss_function(self, expert_outputs, cluster_net_outputs, X, psedo_label):
+        """
+        Compute loss function in a batch
+        Loss = L - Beta * Entropy(cluster_net_outputs) 
+        L = -log [p_i * exp (-(xhat_i - x_i) ** 2)]
+        """
+    
+        # Create one-hot psedo label 
+        # print("Expert output" , expert_outputs)
+        encoded_arr = np.zeros((len(psedo_label), self.n_clusters), dtype=float)
+        for i in range(len(psedo_label)):
+            encoded_arr[i][psedo_label[i]] = 1
+        # print("Cluster network output", cluster_net_outputs)
+        # print("Encoded arr:", encoded_arr)
+        # Cross entropy loss
+        entropy_criterion = nn.CrossEntropyLoss()
+        entropy_loss = entropy_criterion(cluster_net_outputs, torch.tensor(psedo_label, dtype=torch.long))
+        # print("Entropy loss", entropy_loss)
+
+
+        # MOE reconstruction loss
+        loss = 0
+        for i in range(self.n_clusters):
+            mse = -((expert_outputs[i] - X)**2).mean(axis=1)
+            loss += cluster_net_outputs[:, i] * torch.exp(mse)
+
+        moe_loss = -torch.log(loss).sum() 
+        # print('MOE loss', moe_loss)
+        return  moe_loss - self.cl_loss_param * entropy_loss
+        # return moe_loss
+
+
+    def train_one_autoencoder(self, autoencoder, optimizer, criterion, data_loader, number_of_epochs, sparsity, rho, name='main', verbose=False):
+        """
+        Training one autoencoder
+        """
+        print('Training %s ...'%(name))
+        for epoch in range(number_of_epochs):
+
+            running_loss = 0.0
+            autoencoder.train()
+            for batch_index, (data) in enumerate(data_loader):
+                batch_size = data.size()[0]
+                #  Duplicate if batch has one sample (handle one-sample err)
+                if batch_size == 1:
+                    data = torch.cat([data, data], dim=0)
+                    batch_size = 2 
+
+                data = data.to(self.device).view(batch_size, -1)
+                # Get output decoder
+                rec_X = autoencoder(data)
+
+                if sparsity:
+                    # Get latent
+                    reg_loss = criterion(data, rec_X)
+                    sparse_loss = self.sparse_loss(rho=rho, X = data, model=autoencoder)
+                    loss = reg_loss + self.BETA * sparse_loss
+                    # if batch_index & 100 == 0:
+                    #     print("Reg-loss: {} , Sparse-loss: {}".format(reg_loss, sparse_loss))
+                else:
+                    loss = criterion(data, rec_X)
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+                running_loss += loss.data.numpy()
+                if batch_index % 200 ==0 and verbose:
+                    print('epoch %d loss: %.5f batch: %d' % (epoch, running_loss/((batch_index + 1)), (batch_index + 1)*batch_size))
+                if batch_index != 0 and batch_index % 1000 == 0:
+                    break
+        print('Done training %s'%(name))
+
+
+    def pretraining(self, dataloader):
+        """
+        Pretraining step
+        1) Train a single main_autoencoder for the entire dataset
+        2) Apply k-means for the embedding space after training to get label for cluster net
+        3) Training i-th autoencoder using i-th assigned samples by K-means from the entire dataset
+        """
+        #---------Training main_autoencoder---------------
+        criterion = nn.MSELoss()
+        optimizer = torch.optim.Adam(self.main_autoencoder.parameters(), lr=args.lr, weight_decay = args.wd)
+
+        self.train_one_autoencoder(autoencoder=self.main_autoencoder, optimizer=optimizer,
+                                   criterion=criterion, data_loader= dataloader, 
+                                   number_of_epochs=self.pretrain_epochs_main, name= "main_autoencoder", 
+                                   verbose=True,
+                                   sparsity=False,
+                                   rho= self.RHO
+                                   )
+
+        # ----------K-means clustering --------------------------
+        print("------Clustering---------")
+        
+        # Get latent X
+        batch_X = []
+        for batch_idx, (data) in enumerate(dataloader):
+            batch_size = data.size()[0]
+            # Duplicate if batch has one sample
+            if batch_size == 1:
+                data = torch.cat([data, data], dim=0)
+                batch_size = 2 
+            data = data.to(self.device).view(batch_size, -1)
+            latent_X = self.main_autoencoder(data, latent=True)
+            print("BATCH LATENT X", latent_X)
+            batch_X.append(latent_X.detach().cpu().numpy())
+        full_latent_X = np.vstack(batch_X)
+        
+        # Clustering
+        # self.cluster_algo = AgglomerativeClustering(n_clusters=self.n_clusters).fit(full_latent_X)
+        # print("Cluster algo", self.cluster_algo)
+        # self.cluster_algo.fit(full_latent_X)
+        
+        # self.cluster_algo = KMeans(n_clusters=self.n_clusters, n_init= self.n_clusters, init="k-means++", random_state=42).fit(full_latent_X)
+        self.cluster_algo = SpectralClustering(n_clusters=self.n_clusters, random_state=42).fit(full_latent_X)
+        self.psedo_label = self.cluster_algo.labels_
+        print("Done clustering!")
+        print("Original label:", self.psedo_label)
+        # tsne = TSNE(n_components=2, random_state=42)
+        # X_tsne = tsne.fit_transform(full_latent_X)
+        # colors = ['black', 'red']
+        # for i in np.unique(self.cluster_algo.labels_):
+        #     plt.scatter(X_tsne[self.cluster_algo.labels_ == i, 0], X_tsne[self.cluster_algo.labels_ == i, 1], color=colors[i], label=str(i))
+        # plt.xlabel('t-SNE feature 1')
+        # plt.ylabel('t-SNE feature 2')
+        # plt.legend()
+        # plt.show()
+
+        # ---------Training each autoencoder expert with predefined label from K-means---------------
+        for i in range(self.n_clusters):
+            # Get full dataset through batch loop
+            dataset = []
+            for batch_idx, (data) in enumerate(dataloader):
+                batch_size = data.size()[0]
+                # # Duplicate if batch has one sample
+                if batch_size == 1:
+                    data = torch.cat([data, data], dim=0)
+                    batch_size = 2 
+                dataset.append(data.detach().cpu().numpy())
+            dataset = np.vstack(dataset)
+
+            # Extract data for specific expert i
+            data_expert_i = dataset[self.cluster_algo.labels_ == i]
+            data_expert_i = AutoEncoderDataset(data = data_expert_i)
+            dataset_expert_i = DataLoader(data_expert_i, batch_size = args.batch_size, shuffle = False)
+            optimizer = torch.optim.Adam(self.moe[i].parameters(), lr=args.lr, weight_decay = args.wd)
+            criterion = nn.MSELoss()
+            # Train expert_i
+            self.train_one_autoencoder(autoencoder=self.moe[i], optimizer=optimizer,
+                                       criterion=criterion, data_loader=dataset_expert_i, 
+                                       number_of_epochs=self.pretrain_epochs, name="Expert {}".format(i), 
+                                       verbose=True, sparsity=True, rho = self.RHO)
+
+        print("Done Pretraining step !")
+
+        return self.moe, full_latent_X
+    
+
+    def get_expert_outputs(self, X, latent = False):
+        """
+        Get output of experts in a batch
+        Return: List of output of each expert
+        """
+        output = []
+        for i in range(self.n_clusters):
+            if latent:
+                output_expert_i = self.moe[i](X, latent = True)
+            else:
+                output_expert_i = self.moe[i](X)
+            output.append(output_expert_i)
+        return output
+    
+
+    def main_training(self, dataloader, name = "MOE", verbose = True):
+        """
+        Main training to optimize loss function L = -log [p_i * exp (-(xhat_i - x_i) ** 2)]
+        """
+        print('Training %s ...'%(name))
+
+        # Add parameters 
+        params = list(self.moe['cluster_net'].parameters())
+        for i in range(self.n_clusters):
+            params += list(self.moe[i].parameters())        
+            self.moe[i].train()
+            
+        optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay = args.wd)
+
+        for epoch in range(self.main_train_epochs):
+            running_loss = 0.0
+            self.moe['cluster_net'].train()
+            
+            for batch_index, (data) in enumerate(dataloader):
+                batch_size = data.size()[0]
+                # Duplicate if batch has one sample (the last batch)
+                if batch_size == 1:
+                    data = torch.cat([data, data], dim=0)
+                    batch_size = 2 
+                # Get psedo-label
+                psedo_label = self.psedo_label[batch_index: batch_index + batch_size]
+
+                # Get decoder output
+                expert_outputs = self.get_expert_outputs(data)
+                # Get latent output 
+                # latent_outputs = self.get_expert_outputs(data, latent=True)
+                
+                # # Concate k-latent outputs
+                # latent_tensor = latent_outputs[0]
+                # for i in range(1, len(latent_outputs)):
+                #     latent_tensor = torch.hstack((latent_tensor, latent_outputs[i]))
+                # # if batch_index % 100 == 0:
+                # #     # print("Latent tensor", latent_tensor)
+
+                clustering_net_outputs = self.moe['cluster_net'](data)
+
+                # print("Cluster net output", clustering_net_outputs)
+                loss = self.loss_function(expert_outputs=expert_outputs, 
+                                          cluster_net_outputs=clustering_net_outputs, X = data,
+                                          psedo_label=psedo_label)
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+                running_loss += loss.data.numpy()
+                if batch_index % 100 ==0 and verbose:
+                    print('epoch %d loss: %.5f batch: %d' % (epoch, running_loss/((batch_index + 1)), (batch_index + 1)*batch_size))
+                if batch_index != 0 and batch_index % 1000 == 0:
+                    break
+            
+            # Update psedolabel
+            if epoch != 0  and epoch % 10 == 0:
+                self.psedo_label = self.get_final_cluster(dataloader)
+                print("Updated psedo label!")
+                print("######################################")
+                print("New psedo label: ", self.psedo_label)
+                # self.cl_loss_param = self.cl_loss_param * 2 
+        
+        print("Done main training!")
+        return self.moe
+
+
+    def get_final_cluster(self, test_loader):
+        """
+        Assign final cluster for clustering based on cluster_net
+        """
+        # Convert to eval mode
+        for i in range(self.n_clusters):
+            self.moe[i].eval()
+        self.moe['cluster_net'].eval()
+        total_pred = []
+        for batch_idx, (data) in enumerate(test_loader):
+            batch_size = data.size()[0]
+            data = data.view(batch_size, -1).to(self.device)
+            # Get the hard assignment label
+            with torch.no_grad():
+                # # Get latent output 
+                # latent_outputs = self.get_expert_outputs(data, latent=True)
+
+                # # Concate k-latent outputs
+                # latent_tensor = latent_outputs[0]
+                # for i in range(1, len(latent_outputs)):
+                #     latent_tensor = torch.hstack((latent_tensor, latent_outputs[i]))
+
+                cluster_pred = self.moe['cluster_net'](data)
+                cluster_pred = cluster_pred.cpu().numpy()
+                batch_pred = np.argmax(cluster_pred, axis = 1)
+                total_pred.append(batch_pred)
+        total_pred = np.concatenate(total_pred, axis=0)
+        return total_pred
+    
+
+if __name__ == "__main__":
+    # sample_dir = "da_datasets/da_spanish/segments_0.2/1569_[cut_127sec].wav"
+    # dataset = CustomDataset(sample_dir=sample_dir,train= False)
+    # train_loader = DataLoader(dataset, batch_size = args.batch_size, shuffle = False)
+    moe = MoESparseAutoencodersCL(args)
+    # mixture = moe.pretraining(train_loader) 
+    # mixture = moe.main_training(train_loader)
+    # pred = moe.get_final_cluster(train_loader)
+    # print(pred)
+    total_params = sum(p.numel() for p in moe.parameters() if p.requires_grad)
+    print(total_params)
+
+    
+    
+
+    

segment_process.py

@@ -0,0 +1,156 @@
+import os
+import subprocess
+import numpy as np
+import whisper
+import torch
+import argparse
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--data_dir", type=str, default= "./samples", help="data wav dir")
+parser.add_argument("--segment_dir", type=str, default= "./segments_0.4", help="segment dir")
+parser.add_argument("--model_type", type = str, default="tiny", help="model type")
+parser.add_argument("--run_device", type = str, default="cpu", help="run device")
+parser.add_argument("--window_length", type = float, default= 0.4, help="window length")
+parser.add_argument("--overlap", type = float, default= 0, help="overlap")
+
+
+# Define
+opt = parser.parse_args()
+DATA_DIR = opt.data_dir
+SEGMENT_DIR = opt.segment_dir
+model_type = opt.model_type
+run_device = opt.run_device
+window_length = opt.window_length
+overlap = opt.overlap
+
+
+if not os.path.exists(SEGMENT_DIR):
+    os.makedirs(SEGMENT_DIR)
+
+# Load model
+whisper_model = whisper.load_model(model_type, run_device)
+embedding_dims = {"tiny": 384, 'small': 384, 'base': 512, 'medium':1024}
+
+
+
+
+def extract_segment(input_file, output_file,  start_time, end_time):
+    """
+    Extract one segment given start_time and end_time
+    input_file: input .wav file
+    output_file: extracted .wav segment
+    start_time, end_time: start-end time of the segment
+    """
+    # split_file_name = f'./{input_file}_segment_{start_time}_{end_time}.wav'
+    cmd= 'ffmpeg -i '+input_file+' -acodec copy -ss '+str(start_time)+' -to '+str(end_time)+' '+ output_file
+    os.system(cmd)
+
+
+
+def split_audio_with_ffmpeg(input_file, output_dir, segment_length=window_length, overlap=overlap):
+    """
+    Extract all segments from original audio
+    """
+    input_filename = input_file.split("/")[-1]
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    # duration = float(subprocess.check_output(['ffprobe', '-i', input_file, '-show_entries', 'format=duration', '-v', 'quiet', '-of', 'csv=%s' % ("p=0")]))
+    duration = float(subprocess.check_output(['ffprobe', '-i', input_file, '-show_entries', 'format=duration', '-v', 'quiet', '-of', 'csv=%s' % ("p=0")]))
+
+    start_time = 0
+    last_flag = False 
+    while start_time < duration and last_flag == False:
+        end_time = np.round(min(start_time + segment_length, duration), 2)
+        # Cover the last segment
+        if end_time + segment_length > duration:
+            end_time = duration
+            last_flag = True
+        output_file = os.path.join(output_dir, f"{input_filename}_segment_{start_time}_{end_time}.wav")
+        extract_segment(input_file, output_file, start_time, end_time)
+        start_time += segment_length - overlap
+
+
+
+def extract_segment_embedding(segment_dir, save_segment_dir, window_length):
+    """
+    Extract embedding for each segment
+    """
+
+    audio = whisper.load_audio(segment_dir)
+    print("AUDIO SHAPE:", audio.shape)
+
+    # #Duplicate the array to get 30s chunk
+    # audio = np.tile(audio, int(30/window_length))
+
+    # print("AUDIO SHAPE:", audio.shape)
+    mel = whisper.log_mel_spectrogram(audio).to(whisper_model.device)
+
+    
+    # print("MEL SHAPE", mel.shape)
+    #--- this code to create the correct shape of mel spectrogram
+    while True:
+        nF, nT = np.shape(mel)
+        # print(nF, nT)
+        if nT > 3000:
+            mel = mel[:,0:3000]
+            break
+        else:
+            mel = torch.cat((mel, mel), -1)
+    mel = torch.unsqueeze(mel, 0) 
+    
+    wp_emb = whisper_model.embed_audio(mel)
+    print("Wb_emb shape:", wp_emb.shape)
+    # print("WB embedding:", wp_emb)
+   
+    emb_1d  = np.mean(wp_emb.cpu().detach().numpy(), axis=0)
+
+    emb_1d  = np.mean(emb_1d, axis=0)
+
+    emb_1d = np.expand_dims(emb_1d, axis = 0)
+    print("Speaker embedding shape", emb_1d.shape)
+
+    np.save(save_segment_dir + '/{}.npy'.format(segment_dir.split("/")[-1]), emb_1d, allow_pickle=True)
+
+    return emb_1d 
+
+
+
+def delete_segment_after_done(segments_dir):
+    """
+    Delete segment after extracting embedding
+    """
+    for segment in os.listdir(segments_dir):
+        if segment.endswith('.wav'):
+            segment_path = segments_dir + "/" + segment
+            cmd = 'rm '+ segment_path
+            os.system(cmd)
+
+
+
+
+if __name__ == "__main__":
+    data_list = sorted(os.listdir(DATA_DIR))
+
+    for sample in data_list:
+        sample_path = DATA_DIR + "/" + sample
+        segment_save_dir = SEGMENT_DIR + "/" + sample
+        if not os.path.exists(segment_save_dir):
+            os.makedirs(segment_save_dir)
+        # # Extract segments    
+        split_audio_with_ffmpeg(input_file=sample_path, output_dir=segment_save_dir)
+
+        # Extract embedding
+        segment_list = sorted(os.listdir(segment_save_dir), key= lambda x: x.split("_")[-2])
+        for segment in segment_list:
+            segment_path = segment_save_dir + "/" + segment
+            # Extract embedding each segment
+            embed_1d = extract_segment_embedding(segment_dir=segment_path, save_segment_dir= segment_save_dir,window_length=window_length)
+        
+        # Delele segment wav after embeddings are extracted
+        delete_segment_after_done(segments_dir=segment_save_dir)
+    
+
+    
+ 

train_mix_sae.py

@@ -0,0 +1,343 @@
+import torch
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+import numpy as np
+import argparse
+from load_dataset import CustomDataset
+from create_DER import createDER
+from sklearn.cluster import AgglomerativeClustering, KMeans, SpectralClustering, DBSCAN, AffinityPropagation 
+from pyannote.metrics.diarization import DiarizationErrorRate
+from pyannote.core import Segment, Timeline, Annotation
+from sklearn.manifold import TSNE
+from mix_sae import MoESparseAutoencodersCL
+import os
+import pandas as pd
+
+
+######### NOTE: WRITE FILE FOR EACH METHOD ########################
+parser = argparse.ArgumentParser(description='Deep Clustering Network')
+
+# Dataset parameters
+parser.add_argument('--dir', default='./a_dataset/model_size_english/',
+                    help='dataset directory')
+parser.add_argument('--input_dim', type=int, default=1280,
+                    help='input dimension')
+parser.add_argument('--n-classes', type=int, default=2,
+                    help='output dimension')
+
+# Training parameters
+parser.add_argument('--lr', type=float, default=1e-3,
+                    help='learning rate (default: 1e-4)')
+parser.add_argument('--wd', type=float, default=1e-4,
+                    help='weight decay (default: 5e-4)')
+parser.add_argument('--batch-size', type=int, default=16,
+                    help='input batch size for training')
+parser.add_argument('--batch_size_moe', type=int, default=16,
+                    help='input batch size for training')
+
+parser.add_argument('--epoch', type=int, default=50,
+                    help='number of epochs to train')
+parser.add_argument('--pre-epoch', type=int, default=200, 
+                    help='number of pre-train epochs')
+# parser.add_argument('--pretrain_epochs', type=int, default=80,
+#                     help='pretraining step epochs')
+# parser.add_argument('--pretrain', type=bool, default=True,
+#                     help='whether use pre-training')
+# parser.add_argument('--main_train_epochs', type=int, default=150,
+#                     help='main_train epochs')
+# Model parameters
+parser.add_argument('--lamda', type=float, default=1,
+                    help='coefficient of the reconstruction loss')
+parser.add_argument('--beta', type=float, default=0.001, 
+                    help=('coefficient of the regularization term on '
+                            'clustering'))
+parser.add_argument('--hidden-dims', default=[256, 64],
+                    help='learning rate (default: 1e-4)')
+parser.add_argument('--latent_dim', type=int, default=2,
+                    help='latent space dimension')
+parser.add_argument('--n-clusters', type=int, default=2,
+                    help='number of clusters in the latent space')
+
+
+# Utility parameters
+parser.add_argument('--n-jobs', type=int, default=1,
+                    help='number of jobs to run in parallel')
+parser.add_argument('--log-interval', type=int, default=20,
+                    help=('how many batches to wait before logging the '
+                            'training status'))
+parser.add_argument("--window_length", type = float, default= 0.2, help="window length")
+parser.add_argument("--overlap", type = float, default= 0, help="overlap")
+parser.add_argument('--rho', type=float, default=0.2,
+                    help='whether use pre-training')
+parser.add_argument('--pretrain_epochs', type=int, default=10,
+                    help='pretraining step epochs')
+parser.add_argument('--pretrain_epochs_main', type=int, default= 20,
+                    help='pretraining step epochs')
+parser.add_argument('--pretrain', type=bool, default=True,
+                    help='whether use pre-training')
+parser.add_argument('--main_train_epochs', type=int, default =5,
+                    help='main_train epochs')
+parser.add_argument('--sparsity_param', type=float, default=0.01,
+                    help='sparsity constract param')
+parser.add_argument('--cl_loss_param', type=float, default= 1,
+                    help='clasification loss param')
+
+args = parser.parse_args()
+label_dir = args.dir + "/label"
+segment_dir = args.dir + "/large_segments_{}".format(args.window_length)
+pyannote_label_dir = args.dir + "/label"
+window_length = args.window_length
+overlap = args.overlap
+
+sample_list = sorted(os.listdir(segment_dir))
+label_list = sorted(os.listdir(label_dir))
+pyannote_label_list = sorted(os.listdir(pyannote_label_dir))
+
+
+# # Create dataframe to store result
+columns = ["Language", "Filename", "K-means_DER", "K-medoids_DER", "ONLY_PRE_DER" "MOE_CL_DER"]
+# Create dataframe
+df = pd.DataFrame(columns=columns)
+
+for sample, label, py_label in zip(sample_list, label_list, pyannote_label_list):
+    print("Processing segments in folder {}".format(sample))
+    print("Label: ", label)
+    
+    sample_path = segment_dir + "/" + sample
+    label_path = label_dir + '/' + label
+    pyannote_label_path = pyannote_label_dir + "/" + py_label
+    # ## -------------------------- BASELINE ML ----------------------------------
+    # agglo = AgglomerativeClustering(n_clusters=args.n_clusters)
+    # kmeans = KMeans(n_clusters=2)
+
+
+    # # Dataset for ML baseline
+    # ml_dataset = CustomDataset(sample_dir=sample_path, embed_dim=args.input_dim)
+    # data = ml_dataset.data
+    # # reduce_embeddings = dimension_reduce(embeddings=data, reduced_dims=2, reduce_method="pca")
+    # reduce_embeddings = data
+
+    # agglo_res = agglo.fit(reduce_embeddings)
+    # kmeans_res = kmeans.fit(reduce_embeddings)
+
+    
+    # agglo_ref, agglo_hyp = createDER(label_path=label_path, sample_dir=sample_path, prediction=agglo_res.labels_, window_length=window_length, overlap=overlap)
+    # kmean_ref, kmean_hyp = createDER(label_path=label_path, sample_dir=sample_path, prediction=kmeans_res.labels_, window_length=window_length, overlap=overlap)
+    
+    # # print("KMEANS HYP", kmean_hyp)
+    # # print("KMEDOIDS HYP", k_medoids_hyp)
+    # # print("AGGLO HYP", agglo_hyp)
+
+
+
+    # print("K-means label:", kmeans_res.labels_)
+    # print("Agglo label:", agglo_res.labels_)
+
+
+    der = DiarizationErrorRate(collar=0.25, skip_overlap=False)
+    # agglo_error = der(agglo_ref, agglo_hyp)
+    # kmeans_error = der(kmean_ref, kmean_hyp)
+
+
+    # print("Agglo DER: ", agglo_error)
+    # print("K-means DER: ", kmeans_error)
+
+
+    # Check tsne
+    # plt.style.use('grayscale')
+    # check_list = {"k-Means":kmeans_res,"k-Medoids": k_medoids_res, "Agglomerative":agglo_res}
+    # for algo in check_list:
+    #     y_pred = check_list[algo].labels_
+    #     tsne = TSNE(n_components=2, random_state=42)
+    #     X_tsne = tsne.fit_transform(reduce_embeddings)
+    #     colors = ['black', 'aqua']
+    #     for i in np.unique(y_pred):
+    #         plt.scatter(X_tsne[y_pred == i, 0], X_tsne[y_pred == i, 1], color=colors[i], label=str(i))
+    #     plt.xlabel('t-SNE feature 1')
+    #     plt.ylabel('t-SNE feature 2')
+    #     plt.title('t-SNE visualization with cluster labels for {}'.format(algo))
+    #     plt.legend()
+    #     plt.show()
+
+    # # --------------------------- DEEP CLUSTERING ------------------------------------
+
+    train_dataset = CustomDataset(sample_dir=sample_path, embed_dim=args.input_dim)
+    # train_dataset = CustomDataset(sample_dir=sample_path)
+
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=False)
+
+    test_dataset = CustomDataset(sample_dir=sample_path, embed_dim=args.input_dim)
+    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle= False)
+
+    # #Pretrain
+    # rec_loss_list = model.pretrain(train_loader, args.pre_epoch)
+
+    # for e in range(args.epoch):
+    #     model.train()
+    #     model.fit(e, train_loader)
+
+    # y_pred = []
+    # latent_X_list = []
+    # model.eval()  # Set the model to evaluation mode
+
+
+    # for data in test_loader:
+    #     batch_size = data.size()[0]
+    #     data = data.view(batch_size, -1).to(model.device)
+    #     latent_X = model.autoencoder(data, latent=True)
+    #     print('Eval latent x', latent_X)
+    #     latent_X = latent_X.detach().cpu().numpy()
+    #     y_pred.append(model.kmeans.update_assign(latent_X).reshape(-1, 1))
+    #     latent_X_list.append(latent_X)
+
+
+    # y_pred = np.concatenate(y_pred, axis=0)
+    # y_pred = list(np.squeeze(y_pred))
+    # print("Y_pred", y_pred)
+    # latent_X_list = np.concatenate(latent_X_list, axis = 0)
+    # algomerative = KMeans(n_clusters=args.n_clusters)
+    # cluster= algomerative.fit(latent_X_list)
+    # y_pred = cluster.labels_
+    # print("Y_pred", y_pred)
+    # print(latent_X_list)
+    # print(y_pred)
+    # # Perform t-SNE
+    # tsne = TSNE(n_components=2, random_state=42)
+    # X_tsne = tsne.fit_transform(latent_X_list)
+    # colors = ['black', 'aqua']
+    # for i in np.unique(y_pred):
+    #     plt.scatter(X_tsne[y_pred == i, 0], X_tsne[y_pred == i, 1], color=colors[i], label=str(i))
+    # plt.xlabel('t-SNE feature 1')
+    # plt.ylabel('t-SNE feature 2')
+    # plt.title('t-SNE visualization with cluster labels for DCN-SD')
+    # plt.legend()
+    # plt.show()
+
+    # # -------------------SPECTRAL NET-----------------------------
+    # data = torch.tensor(train_dataset.data, dtype=torch.float)
+    # test_data = torch.tensor(test_dataset.data, dtype= torch.float)
+    # spectralnet = SpectralNet(n_clusters=2, 
+    #                           should_use_ae = True,
+    #                           should_use_siamese= False,
+    #                           ae_hiddens = [256, 256, 512, 2],
+    #                           ae_epochs = 150,
+    #                           ae_batch_size = 128, 
+    #                           ae_patience = 30,
+    #                         #   siamese_hiddens = [384, 384, 128, 2],
+    #                         #   siamese_epochs = 150, 
+    #                         #   siamese_batch_size = 128,
+    #                         #   siamese_patience = 30,
+    #                           spectral_hiddens = [384, 384, 512, 2], 
+    #                           spectral_epochs = 300,
+    #                           spectral_batch_size = 128,
+    #                           spectral_patience = 60)
+    # spectralnet.fit(data) # X is the dataset and it should be a torch.Tensor
+    # cluster_assignments = spectralnet.predict(data) # Get the final assignments to cluster
+    # print("Spectral pred", cluster_assignments)
+
+
+    # #------------------MOE ------------------------------------
+    # moe = MoEAutoencoders(args=args)
+    # mixture_moe = moe.pretraining(train_loader) 
+    # mixture_moe = moe.main_training(train_loader)
+    # moe_pred = moe.get_final_cluster(train_loader)
+    # print("MOE pred", moe_pred)
+
+
+    # #------------------MOE SPARITY------------------------------------
+    # moe_spa = MoESparseAutoencoders(args=args)
+    # mixture_moe_spa = moe_spa.pretraining(train_loader) 
+    # mixture_moe_spa = moe_spa.main_training(train_loader)
+    # moe_spa_pred = moe_spa.get_final_cluster(train_loader)
+    # print("MOE SPARITY pred", moe_spa_pred)
+    # plt.style.use("seaborn-v0_8-deep")
+    #------------------MOE SPARITY CL------------------------------------
+    moe_cl = MoESparseAutoencodersCL(args=args)
+    mixture_moe_cl, full_latent_X = moe_cl.pretraining(train_loader)
+    pre_label = moe_cl.psedo_label 
+    # latent_X_list = np.concatenate(full_latent_X, axis = 0)
+    # Perform t-SNE
+    # tsne = TSNE(n_components=2, random_state=42)
+    # X_tsne = tsne.fit_transform(full_latent_X)
+    # colors = ['red', 'green']
+    # for i in np.unique(pre_label):
+    #     plt.scatter(X_tsne[pre_label == i, 0], X_tsne[pre_label == i, 1], color=colors[i], label="Spk {}".format(str(i)))
+    # plt.legend()
+    # # plt.axis("off")
+    # plt.xticks([])
+    # plt.yticks([])
+
+    # plt.show()
+    mixture_moe_cl = moe_cl.main_training(train_loader)
+    moe_cl_pred = moe_cl.get_final_cluster(train_loader)
+    # for data in test_loader:
+    #     batch_size = data.size()[0]
+    #     data = data.view(batch_size, -1).to(model.device)
+    #     latent_X = model.autoencoder(data, latent=True)
+    #     print('Eval latent x', latent_X)
+    #     latent_X = latent_X.detach().cpu().numpy()
+    #     y_pred.append(model.kmeans.update_assign(latent_X).reshape(-1, 1))
+    #     latent_X_list.append(latent_X)
+
+
+    
+
+
+
+    # spt_ref, spt_hyp = createDER(label_path=label_path, sample_dir=sample_path, prediction=cluster_assignments, window_length=window_length, overlap=overlap)
+    # spt_error = der(spt_ref, spt_hyp)
+
+    # print("Spectral Net DER: ", spt_error)
+
+    # moe_ref, moe_hyp = createDER(label_path=label_path, sample_dir=sample_path, prediction=moe_pred, window_length=window_length, overlap=overlap)
+    # moe_error = der(moe_ref, moe_hyp)
+    # print("MOE DER", moe_error)
+
+    # moe_spa_ref, moe_spa_hyp = createDER(label_path=label_path, sample_dir=sample_path, prediction=moe_spa_pred, window_length=window_length, overlap=overlap)
+    # moe_spa_error = der(moe_spa_ref, moe_spa_hyp)
+    # print("MOE SPA DER", moe_spa_error)
+
+    moe_cl_ref, moe_cl_hyp = createDER(label_path=label_path, sample_dir=sample_path, prediction=pre_label, window_length=window_length, overlap=overlap)
+    moe_cl_pre_error = der(moe_cl_ref, moe_cl_hyp)
+
+    print("MOE CL PRE DER", moe_cl_pre_error)
+
+    moe_cl_ref, moe_cl_hyp = createDER(label_path=label_path, sample_dir=sample_path, prediction=moe_cl_pred, window_length=window_length, overlap=overlap)
+    moe_cl_error = der(moe_cl_ref, moe_cl_hyp)
+    print("MOE CL pred", moe_cl_pred)
+    print("MOE CL DER", moe_cl_error)
+
+
+    # # Calculate PYANNOTE DER
+    # ref, py_ref = create_pyannote_timeline(label_path=label_path, pyannote_label_path=pyannote_label_path)
+    # py_error = der(ref, py_ref)
+    # print("Pyannote DER: ", py_error)
+#     # Write result each algo
+#     with open(save_dir + "/{}.txt".format(file_name[:-4]), 'a') as file:
+#         file.write("Algo: {}\n".format(algo))
+#         file.write("Label \n")
+#         file.write(" ".join(label_series) + "\n")
+#         file.write("Pred \n")
+#         file.write(" ".join(pred_series) + "\n")
+#         file.write("DER: {}\n".format(error))
+#         file.write("\n")
+#         file.close()
+
+
+    # Update result to csv
+    language = args.dir.split("/")[-2]
+    file_name = sample
+    new_row = {"Language":language, "Filename": file_name,"MOE_CL":moe_cl_error}
+    
+    df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
+    df.to_csv("1111_co_large_MOE_CL_{}_{}.csv".format(language, window_length), index=False)
+    # Write result each algo
+    with open("./1111_co_large_report_MOE_CL_{}_{}.txt".format(language, window_length), 'a') as file:
+        file.write("Filename: {}\n".format(file_name))
+        file.write("Pred \n")
+        for digit in moe_cl_pred:
+            file.write("{} ".format(digit))
+        file.write("\n")
+        file.write("DER: {}\n".format(moe_cl_error))
+        file.write("\n")
+    file.close()