Upload 2 files

ctcalign.py

@@ -5,39 +5,43 @@ from dataclasses import dataclass
 
 
 
+
 #convert frame-numbers to timestamps in seconds
 # w2v2 step size is about 20ms, or 50 frames per second
 def f2s(fr):
-	return fr/50
-	
-# build labels dict from a processor where it is not directly accessible
-def get_processor_labels(processor,word_sep,max_labels=100):
-	ixs = sorted(list(range(max_labels)),reverse=True)
-	return {processor.tokenizer.decode(n) or word_sep:n for n in ixs}
-	
+    return fr/50
+
+    
 #------------------------------------------
 # setup wav2vec2
 #------------------------------------------
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-torch.random.manual_seed(0)
-max_labels = 100 # any reasonable number higher than vocab + extra + special tokens in any language used
-
 # important to know for CTC decoding - potentially language/model dependent
-model_word_separator = '|'
-model_blank_token = '[PAD]'
+#model_word_separator = '|'
+#model_blank_token = '[PAD]'
+#is_MODEL_PATH="../models/LVL/wav2vec2-large-xlsr-53-icelandic-ep10-1000h"
+
 
-is_MODEL_PATH="carlosdanielhernandezmena/wav2vec2-large-xlsr-53-icelandic-ep10-1000h"
 
+class CTCAligner:
+    
+    def __init__(self, model_path,model_word_separator, model_blank_token):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        torch.random.manual_seed(0)
+        
+        self.model = Wav2Vec2ForCTC.from_pretrained(model_path).to(self.device)
+        self.processor = Wav2Vec2Processor.from_pretrained(model_path)
+
+        # build labels dict from a processor where it is not directly accessible
+        max_labels = 100 # any reasonable number higher than vocab + extra + special tokens in any language used
+        ixs = sorted(list(range(max_labels)),reverse=True)
+        self.labels_dict = {self.processor.tokenizer.decode(n) or model_word_separator:n for n in ixs}
+        
+        self.blank_id = self.labels_dict[model_blank_token]
+        self.model_word_separator = model_word_separator
 
 
 
-model = Wav2Vec2ForCTC.from_pretrained(is_MODEL_PATH).to(device)
-processor = Wav2Vec2Processor.from_pretrained(is_MODEL_PATH)
-labels_dict = get_processor_labels(processor,model_word_separator)
-inverse_dict = {v:k for k,v in labels_dict.items()}
-all_labels = tuple(labels_dict.keys())
-blank_id = labels_dict[model_blank_token]
 
 
 #------------------------------------------
@@ -49,11 +53,11 @@ blank_id = labels_dict[model_blank_token]
 
 # return the label class probability of each audio frame
 # wav is the wav data already read in, NOT the file path.
-def get_frame_probs(wav):
+def get_frame_probs(wav,aligner):
     with torch.inference_mode(): # similar to with torch.no_grad():
-        input_values = processor(wav,sampling_rate=16000).input_values[0]
-        input_values = torch.tensor(input_values, device=device).unsqueeze(0)
-        emits = model(input_values).logits
+        input_values = aligner.processor(wav,sampling_rate=16000).input_values[0]
+        input_values = torch.tensor(input_values, device=aligner.device).unsqueeze(0)
+        emits =  aligner.model(input_values).logits
         emits = torch.log_softmax(emits, dim=-1)
     return emits[0].cpu().detach()
 
@@ -195,28 +199,26 @@ def basic(segs,wsep="|"):
 
 # needs pad labels added to correctly time first segment
 # and therefore add word sep character as placeholder in transcript
-def prep_transcript(xcp):
-    xcp = xcp.replace(' ',model_word_separator)
-    label_ids = [labels_dict[c] for c in xcp]
-    label_ids = [blank_id] + label_ids + [blank_id]
-    xcp = f'{model_word_separator}{xcp}{model_word_separator}'
+def prep_transcript(xcp, aligner):
+    xcp = xcp.replace(' ', aligner.model_word_separator)
+    label_ids = [aligner.labels_dict[c] for c in xcp]
+    label_ids = [aligner.blank_id] +  label_ids + [aligner.blank_id]
+    xcp = f'{ aligner.model_word_separator}{xcp}{aligner.model_word_separator}'
     return xcp,label_ids
-	
-	
-def align(wav_data,transcript):
-	norm_transcript,rec_label_ids = prep_transcript(transcript)
-	emit = get_frame_probs(wav_data)
-	trellis = get_trellis(emit, rec_label_ids, blank_id)
-	path = backtrack(trellis, emit, rec_label_ids, blank_id)
+
+
+
+def align(wav_data,transcript,aligner):
+	norm_transcript,rec_label_ids = prep_transcript(transcript, aligner)
+	emit = get_frame_probs(wav_data,aligner)
+	trellis = get_trellis(emit, rec_label_ids,  aligner.blank_id)
+	path = backtrack(trellis, emit, rec_label_ids,  aligner.blank_id)
 	
 	segments = merge_repeats(path,norm_transcript)
-	words = merge_words(segments, model_word_separator)
+	words = merge_words(segments, aligner.model_word_separator)
 	
 	#segments = [s for s in segments if s[0] != model_word_separator]
 	#return mfalike(segments,words,model_word_separator)
-	return basic(words,model_word_separator), basic(segments,model_word_separator)
-
-
-
+	return basic(words,aligner.model_word_separator), basic(segments,aligner.model_word_separator)
 
 

graph.py

@@ -41,16 +41,15 @@ def get_pitch_tracks(wav_path):
 
 
 
-
 # transcript could be from a corpus with the wav file,
 # input by the user,
 # or from a previous speech recognition process
-def align_and_graph(wav_path, transcript):
+def align_and_graph(wav_path, transcript,lang_aligner):
 
     # fetch data
     #f0_data = get_pitch_tracks(wav_path)
     speech = readwav(wav_path)
-    w_align, seg_align = ctcalign.align(speech,normalise_transcript(transcript))
+    w_align, seg_align = ctcalign.align(speech,normalise_transcript(transcript),lang_aligner)
     
     
     # set up the graph shape
@@ -113,5 +112,3 @@ def align_and_graph(wav_path, transcript):
 
 
 # uppboðssøla bussleiðini viðmerkingar upprunaligur
-
-