batchalign 0.7.14__tar.gz → 0.7.15__tar.gz

{batchalign-0.7.14/batchalign.egg-info → batchalign-0.7.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: batchalign
-Version: 0.7.14
+Version: 0.7.15
 Summary: Python Speech Language Sample Analysis
 Author: Brian MacWhinney, Houjun Liu
 Author-email: macw@cmu.edu, houjun@cmu.edu

{batchalign-0.7.14 → batchalign-0.7.15}/batchalign/models/__init__.py

@@ -1,4 +1,4 @@
-from .utterance import BertUtteranceModel
+from .utterance import BertUtteranceModel, BertCantoneseUtteranceModel
 from .whisper import WhisperASRModel, WhisperFAModel
 from .speaker import NemoSpeakerModel
 from .utils import ASRAudioFile

{batchalign-0.7.14 → batchalign-0.7.15}/batchalign/models/resolve.py

@@ -8,7 +8,7 @@ resolver = {
     "utterance": {
         'eng': "talkbank/CHATUtterance-en",
         "zho": "talkbank/CHATUtterance-zh_CN",
-        "yue": "talkbank/CHATUtterance-zh_CN",
+        "yue": "PolyU-AngelChanLab/Cantonese-Utterance-Segmentation",
     },
     "whisper": {
         'eng': ("talkbank/CHATWhisper-en-large-v1", "openai/whisper-large-v2"),

batchalign-0.7.15/batchalign/models/utterance/__init__.py

@@ -0,0 +1,4 @@
+from .infer import BertUtteranceModel
+from .cantonese_infer import BertCantoneseUtteranceModel
+
+

batchalign-0.7.15/batchalign/models/utterance/cantonese_infer.py

@@ -0,0 +1,164 @@
+import re
+import string
+import random
+
+# tokenization utilities
+import nltk
+from nltk import word_tokenize, sent_tokenize
+
+# torch
+import torch
+from torch.utils.data import dataset 
+from torch.utils.data.dataloader import DataLoader
+from torch.optim import AdamW
+
+# import huggingface utils
+from transformers import AutoTokenizer, BertForTokenClassification
+from transformers import DataCollatorForTokenClassification
+
+# tqdm
+from tqdm import tqdm
+
+# seed device and tokens
+DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+# seed model
+class BertCantoneseUtteranceModel(object):
+
+    def __init__(self, model):
+        # seed tokenizers and model
+        self.tokenizer = AutoTokenizer.from_pretrained(model)
+        self.model = BertForTokenClassification.from_pretrained(model).to(DEVICE)
+        self.max_length = 512
+        self.overlap = 20
+
+        # eval mode
+        self.model.eval()
+        print(f"Model and tokenizer initialized on device: {DEVICE}")
+        print(f"Max length set to {self.max_length} with overlap of {self.overlap}")
+
+    def __call__(self, passage):
+        # Step 1: Clean up passage
+        passage = passage.lower()
+        passage = passage.replace('.','')
+        passage = passage.replace(',','')
+        passage = passage.replace('!','')
+        passage = passage.replace('！','')
+        passage = passage.replace('？','')
+        passage = passage.replace('。','')
+        passage = passage.replace('，','')
+        passage = passage.replace('?','')
+        passage = passage.replace('（','')
+        passage = passage.replace('）','')
+        passage = passage.replace('：','')
+        passage = passage.replace('＊','')
+        passage = passage.replace('ｌ','')
+        
+        
+        # Step 2: Define keywords and split the passage based on them
+        keywords = ['呀', '啦', '喎', '嘞', '㗎喇', '囉', '㗎', '啊', '嗯']  # Replace with your desired keywords
+        
+        chunks = []
+        start = 0
+
+        while start < len(passage):
+            # Find the position of each keyword in the passage starting from the current `start`
+            keyword_positions = [(keyword, passage.find(keyword, start)) for keyword in keywords]
+            # Filter out keywords that are not found (find() returns -1 if not found)
+            keyword_positions = [kp for kp in keyword_positions if kp[1] != -1]
+            
+            if keyword_positions:
+                # Find the keyword that appears first in the passage from current start
+                first_keyword, keyword_pos = min(keyword_positions, key=lambda x: x[1])
+                chunk = passage[start:keyword_pos + len(first_keyword)]
+                chunks.append(chunk)
+                start = keyword_pos + len(first_keyword)
+            else:
+                # No more keywords found, add the rest of the passage as the last chunk
+                chunks.append(passage[start:])
+                break
+        
+        # Debugging: Print number of chunks and their content
+        print(f"Created {len(chunks)} chunks based on keywords.")
+        for i, chunk in enumerate(chunks):
+            print(f"Chunk {i + 1}: {chunk[:100]}...")  # Print the first 100 characters of each chunk
+        
+        # Step 3: Process each chunk and restore punctuation
+        final_passage = []
+        for chunk_index, chunk in enumerate(chunks):
+            print(f"Processing chunk {chunk_index + 1}/{len(chunks)}...")
+
+            # Step 3.1: Split chunk by characters (Chinese tokenization)
+            tokenized_chunk = list(chunk)  # Simply split by characters for Chinese text
+            
+            # Step 3.2: Pass chunk through the tokenizer and model
+            tokd = self.tokenizer.batch_encode_plus([tokenized_chunk],
+                                  return_tensors='pt',
+                                  truncation=True,
+                                  padding=True,
+                                  max_length=self.max_length,
+                                  is_split_into_words=True).to(DEVICE)
+
+            try:
+                # Pass it through the model
+                res = self.model(**tokd).logits
+            except Exception as e:
+                print(f"Error during model inference: {e}")
+                return []
+
+            # Argmax for classification
+            classified_targets = torch.argmax(res, dim=2).cpu()
+
+            # Initialize result tokens list for the current chunk
+            res_toks = []
+            prev_word_idx = None
+
+            # Iterate over tokenized words
+            wids = tokd.word_ids(0)
+            for indx, elem in enumerate(wids):
+                if elem is None or elem == prev_word_idx:
+                    continue
+
+                prev_word_idx = elem
+                action = classified_targets[0][indx]
+
+                # Get the word corresponding to the token
+                w = tokenized_chunk[elem]  # Use tokenized chunk here
+
+                # Fix one word hanging issue (if needed)
+                will_action = False
+                if indx < len(wids) - 2 and classified_targets[0][indx + 1] > 0:
+                    will_action = True
+
+                if not will_action:
+                    # Perform the edits based on model predictions
+                    if action == 1:  # First capital letter
+                        w = w[0].upper() + w[1:]
+                    elif action == 2:  # Add period
+                        w = w + '.'
+                    elif action == 3:  # Add question mark
+                        w = w + '?'
+                    elif action == 4:  # Add exclamation mark
+                        w = w + '!'
+                    elif action == 5:  # Add comma
+                        w = w + ','
+
+                # Append modified word to result list
+                res_toks.append(w)
+
+            # Convert list of tokens back to string and append to final_passage
+            final_passage.append(self.tokenizer.convert_tokens_to_string(res_toks))
+
+        # Step 4: Join processed chunks together into the final passage
+        final_text = ' '.join(final_passage)
+
+        print("Text processing completed. Generating final output...")
+
+        # Optionally, tokenize the final text into sentences based on punctuation
+        try:
+            split_passage = sent_tokenize(final_text)
+        except LookupError:
+            nltk.download('punkt')
+            split_passage = sent_tokenize(final_text)
+
+        return split_passage

{batchalign-0.7.14 → batchalign-0.7.15}/batchalign/models/whisper/infer_asr.py

@@ -33,6 +33,7 @@ import pycountry
 import logging
 L = logging.getLogger("batchalign")
 
+# DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 # DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device("mps") if torch.backends.mps.is_available() else torch.device('cpu')
 # PYTORCH_ENABLE_MPS_FALLBACK=1

{batchalign-0.7.14 → batchalign-0.7.15}/batchalign/pipelines/asr/rev.py

@@ -10,7 +10,7 @@ from batchalign.utils.config import config_read
 
 from batchalign.errors import *
 
-from batchalign.models import BertUtteranceModel, resolve
+from batchalign.models import BertUtteranceModel, BertCantoneseUtteranceModel, resolve
 
 import time
 import pathlib
@@ -49,7 +49,11 @@ class RevEngine(BatchalignEngine):
         self.__client = apiclient.RevAiAPIClient(key)
         if resolve("utterance", lang) != None:
             L.debug("Initializing utterance model...")
-            self.__engine = BertUtteranceModel(resolve("utterance", lang))
+            if lang != "yue":
+                self.__engine = BertUtteranceModel(resolve("utterance", lang))
+            else:
+                # we have special inference procedure for cantonese
+                self.__engine = BertCantoneseUtteranceModel(resolve("utterance", lang))
             L.debug("Done.")
         else:
             self.__engine = None

{batchalign-0.7.14 → batchalign-0.7.15}/batchalign/pipelines/asr/utils.py

@@ -94,7 +94,10 @@ def retokenize_with_engine(intermediate_output, engine):
             tmp = []
 
             for s in new_ut:
-                tmp.append((s, utterance.pop(0)[1]))
+                try:
+                    tmp.append((s, utterance.pop(0)[1]))
+                except IndexError:
+                    continue
 
             final_outputs.append((speaker, tmp+[[delim, [None, None]]]))
 
@@ -159,7 +162,7 @@ def process_generation(output, lang="eng", utterance_engine=None):
                 final_words.append([part.strip(), [cur, cur+div]])
                 cur += div
 
-        lang_2 = pycountry.languages.get(alpha_3=lang).alpha_2
+        lang_2 = "yue" if lang == "yue" else pycountry.languages.get(alpha_3=lang).alpha_2
         def catched_num2words(i):
             if not i.isdigit():
                 return i

{batchalign-0.7.14 → batchalign-0.7.15}/batchalign/pipelines/asr/whisper.py

@@ -1,7 +1,7 @@
 from batchalign.document import *
 from batchalign.pipelines.base import *
 from batchalign.pipelines.asr.utils import *
-from batchalign.models import WhisperASRModel, BertUtteranceModel
+from batchalign.models import WhisperASRModel, BertUtteranceModel, BertCantoneseUtteranceModel
 
 import pycountry
 
@@ -44,7 +44,11 @@ class WhisperEngine(BatchalignEngine):
 
         if resolve("utterance", self.__lang) != None:
             L.debug("Initializing utterance model...")
-            self.__engine = BertUtteranceModel(resolve("utterance", self.__lang))
+            if lang != "yue":
+                self.__engine = BertUtteranceModel(resolve("utterance", lang))
+            else:
+                # we have special inference procedure for cantonese
+                self.__engine = BertCantoneseUtteranceModel(resolve("utterance", lang))
             L.debug("Done.")
         else:
             self.__engine = None

{batchalign-0.7.14 → batchalign-0.7.15}/batchalign/pipelines/fa/wave2vec_fa.py

@@ -154,9 +154,9 @@ class Wave2VecFAEngine(BatchalignEngine):
                 if '\x15' not in ut.text:
                     ut.text = (ut.text+f" \x15{ut.alignment[0]}_{ut.alignment[1]}\x15").strip()
                 else:
-                    ut.text = re.sub("\x15\d+_\d+\x15",
+                    ut.text = re.sub(r"\x15\d+_\d+\x15",
                                      f"\x15{ut.alignment[0]}_{ut.alignment[1]}\x15", ut.text).strip()
             elif ut.text != None:
-                ut.text = re.sub("\x15\d+_\d+\x15", f"", ut.text).strip()
+                ut.text = re.sub(r"\x15\d+_\d+\x15", f"", ut.text).strip()
 
         return doc

{batchalign-0.7.14 → batchalign-0.7.15}/batchalign/pipelines/fa/whisper_fa.py

@@ -179,9 +179,9 @@ class WhisperFAEngine(BatchalignEngine):
                 if '\x15' not in ut.text:
                     ut.text = (ut.text+f" \x15{ut.alignment[0]}_{ut.alignment[1]}\x15").strip()
                 else:
-                    ut.text = re.sub("\x15\d+_\d+\x15",
+                    ut.text = re.sub(r"\x15\d+_\d+\x15",
                                      f"\x15{ut.alignment[0]}_{ut.alignment[1]}\x15", ut.text).strip()
             elif ut.text != None:
-                ut.text = re.sub("\x15\d+_\d+\x15", f"", ut.text).strip()
+                ut.text = re.sub(r"\x15\d+_\d+\x15", f"", ut.text).strip()
 
         return doc

batchalign-0.7.15/batchalign/version

@@ -0,0 +1,3 @@
+0.7.15
+Feburary 23rd, 2025
+Whisper ASR with Cantonese and tokenization!

{batchalign-0.7.14 → batchalign-0.7.15/batchalign.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: batchalign
-Version: 0.7.14
+Version: 0.7.15
 Summary: Python Speech Language Sample Analysis
 Author: Brian MacWhinney, Houjun Liu
 Author-email: macw@cmu.edu, houjun@cmu.edu

{batchalign-0.7.14 → batchalign-0.7.15}/batchalign.egg-info/SOURCES.txt

@@ -40,6 +40,7 @@ batchalign/models/training/__init__.py
 batchalign/models/training/run.py
 batchalign/models/training/utils.py
 batchalign/models/utterance/__init__.py
+batchalign/models/utterance/cantonese_infer.py
 batchalign/models/utterance/dataset.py
 batchalign/models/utterance/execute.py
 batchalign/models/utterance/infer.py

batchalign-0.7.14/batchalign/models/utterance/__init__.py

@@ -1,2 +0,0 @@
-from .infer import BertUtteranceModel
-

batchalign-0.7.14/batchalign/version

@@ -1,3 +0,0 @@
-0.7.14
-Feburary 19nd, 2025
-machine translation!