Uploading ngram base model

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,3 +1,142 @@
----
-license: cc-by-4.0
----
+# llm_speaker_tagging
+
+SLT 2024 Challenge: Post-ASR-Speaker-Tagging Baseline
+
+# Project Name
+
+SLT 2024 Challenge GenSEC Track 2: Post-ASR-Speaker-Tagging Baseline 
+
+## Features
+
+- Data download and cleaning
+- n-gram + beam search decoder based baselinee system 
+
+## Installation
+
+Run the following commands at the main level of this repository.
+
+### Conda Environment
+
+```
+conda create --name llmspk python=3.10
+```
+### Install requirements
+
+You need to install the following packages 
+
+```
+kenlm
+arpa
+numpy
+hydra-core
+meeteval
+tqdm
+requests
+simplejson
+pydiardecode @ git+https://github.com/tango4j/pydiardecode@main
+```
+
+Simply install all the requirments. 
+
+```
+pip install -r requirements.txt
+```
+
+### Download ARPA language model
+
+```
+mkdir -p arpa_model
+cd arpa_model
+wget https://kaldi-asr.org/models/5/4gram_small.arpa.gz
+gunzip 4gram_small.arpa.gz
+```
+
+### Download track-2 challenge dev set and eval set 
+
+Clone the dataset from Hugging Face server.
+```
+git clone https://huggingface.co/datasets/GenSEC-LLM/SLT-Task2-Post-ASR-Speaker-Tagging
+```
+
+```
+find . $PWD/SLT-Task2-Post-ASR-Speaker-Tagging/err_source_text/dev -name *.seglst.json > err_dev.src.list
+find . $PWD/SLT-Task2-Post-ASR-Speaker-Tagging/ref_annotated_text/dev -name *.seglst.json > err_dev.ref.list
+```
+
+### Launch the baseline script
+
+Now you are ready to launch the script.
+Launch the baseline script `run_speaker_tagging_beam_search.sh`
+
+```
+BASEPATH=${PWD}
+DIAR_LM_PATH=$BASEPATH/arpa_model/4gram_small.arpa
+ASRDIAR_FILE_NAME=err_dev
+WORKSPACE=$BASEPATH/SLT-Task2-Post-ASR-Speaker-Tagging
+INPUT_ERROR_SRC_LIST_PATH=$BASEPATH/$ASRDIAR_FILE_NAME.src.list
+GROUNDTRUTH_REF_LIST_PATH=$BASEPATH/$ASRDIAR_FILE_NAME.ref.list
+DIAR_OUT_DOWNLOAD=$WORKSPACE/short2_all_seglst_infer
+mkdir -p $DIAR_OUT_DOWNLOAD
+
+### SLT 2024 Speaker Tagging Setting v1.0.2
+ALPHA=0.4
+BETA=0.04
+PARALLEL_CHUNK_WORD_LEN=100
+BEAM_WIDTH=16
+WORD_WINDOW=32
+PEAK_PROB=0.95
+USE_NGRAM=True
+LM_METHOD=ngram
+
+# Get the base name of the test_manifest and remove extension
+UNIQ_MEMO=$(basename "${INPUT_ERROR_SRC_LIST_PATH}" .json | sed 's/\./_/g') 
+echo "UNIQ MEMO:" $UNIQ_MEMO
+TRIAL=telephonic
+BATCH_SIZE=11
+
+rm $WORKSPACE/$ASRDIAR_FILE_NAME.src.seglst.json
+rm $WORKSPACE/$ASRDIAR_FILE_NAME.ref.seglst.json
+rm $WORKSPACE/$ASRDIAR_FILE_NAME.hyp.seglst.json
+
+python $BASEPATH/speaker_tagging_beamsearch.py \
+    port=[5501,5502,5511,5512,5521,5522,5531,5532] \
+    arpa_language_model=$DIAR_LM_PATH \
+    batch_size=$BATCH_SIZE \
+    groundtruth_ref_list_path=$GROUNDTRUTH_REF_LIST_PATH \
+    input_error_src_list_path=$INPUT_ERROR_SRC_LIST_PATH \
+    parallel_chunk_word_len=$PARALLEL_CHUNK_WORD_LEN \
+    use_ngram=$USE_NGRAM \
+    alpha=$ALPHA \
+    beta=$BETA \
+    beam_width=$BEAM_WIDTH \
+    word_window=$WORD_WINDOW \
+    peak_prob=$PEAK_PROB \
+    out_dir=$DIAR_OUT_DOWNLOAD 
+```
+
+### Evaluate 
+
+We use [MeetEval](https://github.com/fgnt/meeteval) software to evaluate `cpWER`.
+cpWER measures both speaker tagging and word error rate (WER) by testing all the permutation of trancripts and choosing the permutation that 
+gives the lowest error.
+
+```
+echo "Evaluating the original source transcript."
+meeteval-wer cpwer -h $WORKSPACE/$ASRDIAR_FILE_NAME.src.seglst.json -r $WORKSPACE/$ASRDIAR_FILE_NAME.ref.seglst.json 
+echo "Source     cpWER: " $(jq '.error_rate' "[ $WORKSPACE/$ASRDIAR_FILE_NAME.src.seglst_cpwer.json) ]"
+
+echo "Evaluating the original hypothesis transcript."
+meeteval-wer cpwer -h $WORKSPACE/$ASRDIAR_FILE_NAME.hyp.seglst.json -r $WORKSPACE/$ASRDIAR_FILE_NAME.ref.seglst.json 
+echo "Hypothesis cpWER: " $(jq '.error_rate'  $WORKSPACE/$ASRDIAR_FILE_NAME.hyp.seglst_cpwer.json)
+```
+
+### Reference
+
+@inproceedings{park2024enhancing,
+  title={Enhancing speaker diarization with large language models: A contextual beam search approach},
+  author={Park, Tae Jin and Dhawan, Kunal and Koluguri, Nithin and Balam, Jagadeesh},
+  booktitle={ICASSP 2024-2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
+  pages={10861--10865},
+  year={2024},
+  organization={IEEE}
+}

beam_search_utils.py

@@ -0,0 +1,325 @@
+from tqdm import tqdm
+from typing import Dict, List
+from pydiardecode import build_diardecoder
+import numpy as np
+import copy 
+import os
+import json 
+import concurrent.futures
+import kenlm
+
+__INFO_TAG__ = "[INFO]"
+
+class SpeakerTaggingBeamSearchDecoder:
+    def __init__(self, loaded_kenlm_model: kenlm, cfg: dict):
+        self.realigning_lm_params = cfg
+        self.realigning_lm = self._load_realigning_LM(loaded_kenlm_model=loaded_kenlm_model)
+        self._SPLITSYM = "@"
+
+    def _load_realigning_LM(self, loaded_kenlm_model: kenlm):
+        """
+        Load ARPA language model for realigning speaker labels for words.
+        """
+        diar_decoder = build_diardecoder(
+            loaded_kenlm_model=loaded_kenlm_model,
+            kenlm_model_path=self.realigning_lm_params['arpa_language_model'], 
+            alpha=self.realigning_lm_params['alpha'], 
+            beta=self.realigning_lm_params['beta'],
+            word_window=self.realigning_lm_params['word_window'],
+            use_ngram=self.realigning_lm_params['use_ngram'],
+        )
+        return diar_decoder
+
+    def realign_words_with_lm(self, word_dict_seq_list: List[Dict[str, float]], speaker_count: int = None, port_num=None) -> List[Dict[str, float]]:
+        if speaker_count is None:
+            spk_list = []
+            for k, line_dict in enumerate(word_dict_seq_list):
+                _, spk_label = line_dict['word'], line_dict['speaker']
+                spk_list.append(spk_label)
+        else:
+            spk_list = [ f"speaker_{k}" for k in range(speaker_count)]
+
+        realigned_list = self.realigning_lm.decode_beams(beam_width=self.realigning_lm_params['beam_width'],
+                                                         speaker_list=sorted(list(set(spk_list))), 
+                                                         word_dict_seq_list=word_dict_seq_list,
+                                                         port_num=port_num)
+        return realigned_list
+
+    def beam_search_diarization(
+        self,
+        trans_info_dict: Dict[str, Dict[str, list]],
+        port_num: List[int] = None,
+    ) -> Dict[str, Dict[str, float]]:
+        """
+        Match the diarization result with the ASR output.
+        The words and the timestamps for the corresponding words are matched in a for loop.
+
+        Args:
+
+        Returns:
+            trans_info_dict (dict):
+                Dictionary containing word timestamps, speaker labels and words from all sessions.
+                Each session is indexed by a unique ID.
+        """
+        for uniq_id, session_dict in tqdm(trans_info_dict.items(), total=len(trans_info_dict), disable=True):
+            word_dict_seq_list = session_dict['words']
+            output_beams = self.realign_words_with_lm(word_dict_seq_list=word_dict_seq_list, speaker_count=session_dict['speaker_count'], port_num=port_num)
+            word_dict_seq_list = output_beams[0][2]
+            trans_info_dict[uniq_id]['words'] = word_dict_seq_list
+        return trans_info_dict
+
+    def merge_div_inputs(self, div_trans_info_dict, org_trans_info_dict, win_len=250, word_window=16):
+        """
+        Merge the outputs of parallel processing.
+        """
+        uniq_id_list = list(org_trans_info_dict.keys())
+        sub_div_dict = {}
+        for seq_id in div_trans_info_dict.keys():
+            div_info = seq_id.split(self._SPLITSYM)
+            uniq_id, sub_idx, total_count = div_info[0], int(div_info[1]), int(div_info[2])
+            if uniq_id not in sub_div_dict:
+                sub_div_dict[uniq_id] = [None] * total_count
+            sub_div_dict[uniq_id][sub_idx] = div_trans_info_dict[seq_id]['words']
+                
+        for uniq_id in uniq_id_list:
+            org_trans_info_dict[uniq_id]['words'] = []
+            for k, div_words in enumerate(sub_div_dict[uniq_id]):
+                if k == 0:
+                    div_words = div_words[:win_len]
+                else:
+                    div_words = div_words[word_window:]
+                org_trans_info_dict[uniq_id]['words'].extend(div_words)
+        return org_trans_info_dict
+    
+    def divide_chunks(self, trans_info_dict, win_len, word_window, port):
+        """
+        Divide word sequence into chunks of length `win_len` for parallel processing.    
+
+        Args:
+            trans_info_dict (_type_): _description_
+            diar_logits (_type_): _description_
+            win_len (int, optional): _description_. Defaults to 250.
+        """
+        if len(port) > 1:
+            num_workers = len(port) 
+        else:
+            num_workers = 1
+        div_trans_info_dict = {}
+        for uniq_id in trans_info_dict.keys():
+            uniq_trans = trans_info_dict[uniq_id]
+            del uniq_trans['status']
+            del uniq_trans['transcription']
+            del uniq_trans['sentences']
+            word_seq = uniq_trans['words']
+
+            div_word_seq = [] 
+            if win_len is None:
+                win_len = int(np.ceil(len(word_seq)/num_workers))
+            n_chunks = int(np.ceil(len(word_seq)/win_len))
+            
+            for k in range(n_chunks):
+                div_word_seq.append(word_seq[max(k*win_len - word_window, 0):(k+1)*win_len])
+            
+            total_count = len(div_word_seq)
+            for k, w_seq in enumerate(div_word_seq):
+                seq_id = uniq_id + f"{self._SPLITSYM}{k}{self._SPLITSYM}{total_count}"
+                div_trans_info_dict[seq_id] = dict(uniq_trans)
+                div_trans_info_dict[seq_id]['words'] = w_seq
+        return div_trans_info_dict
+
+
+def run_mp_beam_search_decoding(
+    speaker_beam_search_decoder, 
+    loaded_kenlm_model, 
+    trans_info_dict, 
+    org_trans_info_dict, 
+    div_mp, 
+    win_len, 
+    word_window, 
+    port=None, 
+    use_ngram=False
+    ):
+    if len(port) > 1:
+        port = [int(p) for p in port]
+    if use_ngram:
+        port = [None]
+        num_workers = 36
+    else:
+        num_workers = len(port)
+    
+    uniq_id_list = sorted(list(trans_info_dict.keys() ))
+    tp = concurrent.futures.ProcessPoolExecutor(max_workers=num_workers)
+    futures = []
+
+    count = 0
+    for uniq_id in uniq_id_list:
+        print(f"{__INFO_TAG__} Running beam search decoding for {uniq_id}...")
+        if port is not None:
+            port_num = port[count % len(port)]    
+        else:
+            port_num = None
+        count += 1
+        uniq_trans_info_dict = {uniq_id: trans_info_dict[uniq_id]}
+        futures.append(tp.submit(speaker_beam_search_decoder.beam_search_diarization, uniq_trans_info_dict, port_num=port_num))
+
+    pbar = tqdm(total=len(uniq_id_list), desc="Running beam search decoding", unit="files")
+    count = 0
+    output_trans_info_dict = {}
+    for done_future in concurrent.futures.as_completed(futures):
+        count += 1
+        pbar.update()
+        output_trans_info_dict.update(done_future.result())
+    pbar.close() 
+    tp.shutdown()
+    if div_mp:
+        output_trans_info_dict = speaker_beam_search_decoder.merge_div_inputs(div_trans_info_dict=output_trans_info_dict, 
+                                                                              org_trans_info_dict=org_trans_info_dict, 
+                                                                              win_len=win_len, 
+                                                                              word_window=word_window)
+    return output_trans_info_dict
+
+def count_num_of_spks(json_trans_list):
+    spk_set = set()
+    for sentence_dict in json_trans_list:
+        spk_set.add(sentence_dict['speaker'])
+    speaker_map = { spk_str: idx for idx, spk_str in enumerate(spk_set)}
+    return speaker_map
+
+def add_placeholder_speaker_softmax(json_trans_list, peak_prob=0.94 ,max_spks=4): 
+    nemo_json_dict = {}
+    word_dict_seq_list = []
+    if peak_prob > 1 or peak_prob < 0:
+        raise ValueError(f"peak_prob must be between 0 and 1 but got {peak_prob}")
+    speaker_map = count_num_of_spks(json_trans_list)
+    base_array = np.ones(max_spks) * (1 - peak_prob)/(max_spks-1)
+    stt_sec, end_sec = None, None
+    for sentence_dict in json_trans_list:
+        word_list = sentence_dict['words'].split()
+        speaker = sentence_dict['speaker']
+        for word in word_list:
+            speaker_softmax = copy.deepcopy(base_array)
+            speaker_softmax[speaker_map[speaker]] = peak_prob
+            word_dict_seq_list.append({'word': word, 
+                                    'start_time': stt_sec, 
+                                    'end_time': end_sec, 
+                                    'speaker': speaker_map[speaker], 
+                                    'speaker_softmax': speaker_softmax}
+                                    )
+    nemo_json_dict.update({'words': word_dict_seq_list, 
+                           'status': "success", 
+                           'sentences': json_trans_list, 
+                           'speaker_count': len(speaker_map), 
+                           'transcription': None}
+                        )
+    return nemo_json_dict
+
+def convert_nemo_json_to_seglst(trans_info_dict):
+    seglst_seq_list = []
+    seg_lst_dict, spk_wise_trans_sessions = {}, {}
+    for uniq_id in trans_info_dict.keys():
+        spk_wise_trans_sessions[uniq_id] = {}
+        seglst_seq_list = []
+        word_seq_list = trans_info_dict[uniq_id]['words']
+        prev_speaker, sentence = None, ''
+        for widx, word_dict in enumerate(word_seq_list):
+            curr_speaker = word_dict['speaker']
+
+            # For making speaker wise transcriptions
+            word = word_dict['word']
+            if curr_speaker not in spk_wise_trans_sessions[uniq_id]:
+                spk_wise_trans_sessions[uniq_id][curr_speaker] = word
+            elif curr_speaker in spk_wise_trans_sessions[uniq_id]:
+                spk_wise_trans_sessions[uniq_id][curr_speaker] = f"{spk_wise_trans_sessions[uniq_id][curr_speaker]} {word_dict['word']}"
+
+            # For making segment wise transcriptions
+            if curr_speaker!= prev_speaker and prev_speaker is not None:
+                seglst_seq_list.append({'session_id': uniq_id, 
+                                        'words': sentence.strip(), 
+                                        'start_time': 0.0,
+                                        'end_time': 0.0,
+                                        'speaker': prev_speaker, 
+                })
+                sentence = word_dict['word']
+            else:
+                sentence = f"{sentence} {word_dict['word']}"
+            prev_speaker = curr_speaker
+
+        # For the last word:
+        # (1) If there is no speaker change, add the existing sentence and exit the loop
+        # (2) If there is a speaker change, add the last word and exit the loop
+        if widx == len(word_seq_list) - 1:
+            seglst_seq_list.append({'session_id': uniq_id, 
+                                    'words': sentence.strip(), 
+                                    'start_time': 0.0,
+                                    'end_time': 0.0,
+                                    'speaker': curr_speaker, 
+            })
+        seg_lst_dict[uniq_id] = seglst_seq_list
+    return seg_lst_dict
+
+def load_input_jsons(input_error_src_list_path, ext_str=".seglst.json", peak_prob=0.94, max_spks=4):
+    trans_info_dict = {}
+    json_filepath_list = open(input_error_src_list_path).readlines()
+    for json_path in json_filepath_list:
+        json_path = json_path.strip()
+        uniq_id = os.path.split(json_path)[-1].split(ext_str)[0]
+        if os.path.exists(json_path):
+            with open(json_path, "r") as file:
+                json_trans = json.load(file)
+        else:
+            raise FileNotFoundError(f"{json_path} does not exist. Aborting.")
+        nemo_json_dict = add_placeholder_speaker_softmax(json_trans, peak_prob=peak_prob, max_spks=max_spks)
+        trans_info_dict[uniq_id] = nemo_json_dict
+    return trans_info_dict
+
+def load_reference_jsons(reference_seglst_list_path,  ext_str=".seglst.json"):
+    reference_info_dict = {}
+    json_filepath_list = open(reference_seglst_list_path).readlines()
+    for json_path in json_filepath_list:
+        json_path = json_path.strip()
+        uniq_id = os.path.split(json_path)[-1].split(ext_str)[0]
+        if os.path.exists(json_path):
+            with open(json_path, "r") as file:
+                json_trans = json.load(file)
+        else:
+            raise FileNotFoundError(f"{json_path} does not exist. Aborting.")
+        json_trans_uniq_id = []
+        for sentence_dict in json_trans:
+            sentence_dict['session_id'] = uniq_id
+            json_trans_uniq_id.append(sentence_dict)
+        reference_info_dict[uniq_id] = json_trans_uniq_id
+    return reference_info_dict 
+
+def write_seglst_jsons(
+    seg_lst_sessions_dict: dict, 
+    input_error_src_list_path: str, 
+    diar_out_path: str, 
+    ext_str: str, 
+    write_individual_seglst_jsons=True
+    ):
+    """
+    Writes the segment list (seglst) JSON files to the output directory.
+
+    Parameters:
+        seg_lst_sessions_dict (dict): A dictionary containing session IDs as keys and their corresponding segment lists as values.
+        input_error_src_list_path (str): The path to the input error source list file.
+        diar_out_path (str): The path to the output directory where the seglst JSON files will be written.
+        type_string (str): A string representing the type of the seglst JSON files (e.g., 'hyp' for hypothesis or 'ef' for reference).
+        write_individual_seglst_jsons (bool, optional): A flag indicating whether to write individual seglst JSON files for each session. Defaults to True.
+
+    Returns:
+        None
+    """
+    total_infer_list = []
+    total_output_filename = os.path.split(input_error_src_list_path)[-1].replace(".list", "")
+    for session_id, seg_lst_list in seg_lst_sessions_dict.items():
+        total_infer_list.extend(seg_lst_list)
+        if write_individual_seglst_jsons:
+            print(f"{__INFO_TAG__} Writing {diar_out_path}/{session_id}.seglst.json")
+            with open(f'{diar_out_path}/{session_id}.seglst.json', 'w') as file:
+                json.dump(seg_lst_list, file, indent=4)  # indent=4 for pretty printing
+
+    print(f"{__INFO_TAG__} Writing {diar_out_path}/{session_id}.seglst.json")
+    total_output_filename = total_output_filename.replace("src", ext_str).replace("ref", ext_str)
+    with open(f'{diar_out_path}/../{total_output_filename}.seglst.json', 'w') as file:
+        json.dump(total_infer_list, file, indent=4)  # indent=4 for pretty printing

requirements.txt

@@ -0,0 +1,9 @@
+kenlm
+arpa
+numpy
+hydra-core
+meeteval
+tqdm
+requests
+simplejson
+pydiardecode @ git+https://github.com/tango4j/pydiardecode@main

run_speaker_tagging_beam_search.sh

@@ -0,0 +1,59 @@
+
+### Speaker Tagging Task-2 Parameters
+
+
+BASEPATH=${PWD}
+DIAR_LM_PATH=$BASEPATH/arpa_model/4gram_small.arpa
+ASRDIAR_FILE_NAME=err_dev
+WORKSPACE=$BASEPATH/SLT-Task2-Post-ASR-Speaker-Tagging
+INPUT_ERROR_SRC_LIST_PATH=$BASEPATH/$ASRDIAR_FILE_NAME.src.list
+GROUNDTRUTH_REF_LIST_PATH=$BASEPATH/$ASRDIAR_FILE_NAME.ref.list
+DIAR_OUT_DOWNLOAD=$WORKSPACE/short2_all_seglst_infer
+mkdir -p $DIAR_OUT_DOWNLOAD
+
+
+### SLT 2024 Speaker Tagging Setting v1.0.2
+ALPHA=0.4
+BETA=0.04
+PARALLEL_CHUNK_WORD_LEN=100
+BEAM_WIDTH=16
+WORD_WINDOW=32
+PEAK_PROB=0.95
+USE_NGRAM=True
+LM_METHOD=ngram
+
+# Get the base name of the test_manifest and remove extension
+UNIQ_MEMO=$(basename "${INPUT_ERROR_SRC_LIST_PATH}" .json | sed 's/\./_/g') 
+echo "UNIQ MEMO:" $UNIQ_MEMO
+TRIAL=telephonic
+BATCH_SIZE=11
+
+
+rm $WORKSPACE/$ASRDIAR_FILE_NAME.src.seglst.json
+rm $WORKSPACE/$ASRDIAR_FILE_NAME.ref.seglst.json
+rm $WORKSPACE/$ASRDIAR_FILE_NAME.hyp.seglst.json
+
+
+python $BASEPATH/speaker_tagging_beamsearch.py \
+    port=[5501,5502,5511,5512,5521,5522,5531,5532] \
+    arpa_language_model=$DIAR_LM_PATH \
+    batch_size=$BATCH_SIZE \
+    groundtruth_ref_list_path=$GROUNDTRUTH_REF_LIST_PATH \
+    input_error_src_list_path=$INPUT_ERROR_SRC_LIST_PATH \
+    parallel_chunk_word_len=$PARALLEL_CHUNK_WORD_LEN \
+    use_ngram=$USE_NGRAM \
+    alpha=$ALPHA \
+    beta=$BETA \
+    beam_width=$BEAM_WIDTH \
+    word_window=$WORD_WINDOW \
+    peak_prob=$PEAK_PROB \
+    out_dir=$DIAR_OUT_DOWNLOAD 
+
+
+echo "Evaluating the original source transcript."
+meeteval-wer cpwer -h $WORKSPACE/$ASRDIAR_FILE_NAME.src.seglst.json -r $WORKSPACE/$ASRDIAR_FILE_NAME.ref.seglst.json 
+echo "Source     cpWER: " $(jq '.error_rate' "[ $WORKSPACE/$ASRDIAR_FILE_NAME.src.seglst_cpwer.json) ]"
+
+echo "Evaluating the original hypothesis transcript."
+meeteval-wer cpwer -h $WORKSPACE/$ASRDIAR_FILE_NAME.hyp.seglst.json -r $WORKSPACE/$ASRDIAR_FILE_NAME.ref.seglst.json 
+echo "Hypothesis cpWER: " $(jq '.error_rate'  $WORKSPACE/$ASRDIAR_FILE_NAME.hyp.seglst_cpwer.json)

speaker_tagging_beamsearch.py

@@ -0,0 +1,76 @@
+import hydra
+from typing import List, Optional
+from dataclasses import dataclass, field
+import kenlm
+from beam_search_utils import (
+    SpeakerTaggingBeamSearchDecoder,
+    load_input_jsons,
+    load_reference_jsons,
+    write_seglst_jsons,
+    run_mp_beam_search_decoding,
+    convert_nemo_json_to_seglst,
+)
+from hydra.core.config_store import ConfigStore
+
+__INFO_TAG__ = "[INFO]"
+
+@dataclass
+class RealigningLanguageModelParameters:
+    batch_size: int = 32
+    use_mp: bool = True
+    input_error_src_list_path: Optional[str] = None
+    groundtruth_ref_list_path: Optional[str] = None
+    arpa_language_model: Optional[str] = None
+    word_window: int = 32
+    port: List[int] = field(default_factory=list)
+    parallel_chunk_word_len: int = 250
+    use_ngram: bool = True
+    peak_prob: float = 0.95
+    alpha: float = 0.5
+    beta: float = 0.05
+    beam_width: int = 16
+    out_dir: Optional[str] = None
+
+cs = ConfigStore.instance()
+cs.store(name="config", node=RealigningLanguageModelParameters)
+
+@hydra.main(config_name="config", version_base="1.1")
+def main(cfg: RealigningLanguageModelParameters) -> None:
+    trans_info_dict = load_input_jsons(input_error_src_list_path=cfg.input_error_src_list_path, peak_prob=float(cfg.peak_prob))
+    reference_info_dict  = load_reference_jsons(reference_seglst_list_path=cfg.groundtruth_ref_list_path)
+    source_info_dict = load_reference_jsons(reference_seglst_list_path=cfg.input_error_src_list_path)
+    loaded_kenlm_model = kenlm.Model(cfg.arpa_language_model)
+    
+    speaker_beam_search_decoder = SpeakerTaggingBeamSearchDecoder(loaded_kenlm_model=loaded_kenlm_model, cfg=cfg)
+    
+    div_trans_info_dict = speaker_beam_search_decoder.divide_chunks(trans_info_dict=trans_info_dict, 
+                                                                    win_len=cfg.parallel_chunk_word_len, 
+                                                                    word_window=cfg.word_window,
+                                                                    port=cfg.port,)
+    
+    trans_info_dict = run_mp_beam_search_decoding(speaker_beam_search_decoder, 
+                                                    loaded_kenlm_model=loaded_kenlm_model,
+                                                    trans_info_dict=div_trans_info_dict, 
+                                                    org_trans_info_dict=trans_info_dict, 
+                                                    div_mp=True,
+                                                    win_len=cfg.parallel_chunk_word_len,
+                                                    word_window=cfg.word_window,
+                                                    port=cfg.port,
+                                                    use_ngram=cfg.use_ngram,
+                                                    )
+    hypothesis_sessions_dict = convert_nemo_json_to_seglst(trans_info_dict) 
+    
+    write_seglst_jsons(hypothesis_sessions_dict, input_error_src_list_path=cfg.input_error_src_list_path, diar_out_path=cfg.out_dir, ext_str='hyp')
+    write_seglst_jsons(reference_info_dict, input_error_src_list_path=cfg.groundtruth_ref_list_path, diar_out_path=cfg.out_dir, ext_str='ref')
+    write_seglst_jsons(source_info_dict, input_error_src_list_path=cfg.groundtruth_ref_list_path, diar_out_path=cfg.out_dir, ext_str='src')
+    print(f"{__INFO_TAG__} Parameters used: \
+            \n ALPHA: {cfg.alpha} \
+            \n BETA: {cfg.beta} \
+            \n BEAM WIDTH: {cfg.beam_width} \
+            \n Word Window: {cfg.word_window} \
+            \n Use Ngram: {cfg.use_ngram} \
+            \n Chunk Word Len: {cfg.parallel_chunk_word_len} \
+            \n SpeakerLM Model: {cfg.arpa_language_model}") \
+
+if __name__ == '__main__':
+    main()