Update app.py

app.py

@@ -0,0 +1,144 @@
+from os.path import join
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC, Wav2Vec2CTCTokenizer
+import torch
+import pandas as pd
+import librosa
+import gradio as gr
+from gradio.components import Audio, Dropdown, Textbox
+
+
+Attributes = {'Dental':2,
+              'Labial':4,
+             'Consonant':15,
+             'Vowel':19,
+             'Fricative':21,
+             'Nasal':22,
+             'Stop':23,
+             'Affricate':25,
+             'Voiced':31,
+             'Bilabial':32,
+             }
+
+#define groups
+#make sure that all phonemes covered in each group
+g1 = ['p_alveolar','n_alveolar']
+g2 = ['p_palatal','n_palatal']
+g3 = ['p_dental','n_dental']
+g4 = ['p_glottal','n_glottal']
+g5 = ['p_labial','n_labial']
+g6 = ['p_velar','n_velar']
+g7 = ['p_anterior','n_anterior']
+g8 = ['p_posterior','n_posterior']
+g9 = ['p_retroflex','n_retroflex']
+g10 = ['p_mid','n_mid']
+g11 = ['p_high_v','n_high_v']
+g12 = ['p_low','n_low']
+g13 = ['p_front','n_front']
+g14 = ['p_back','n_back']
+g15 = ['p_central','n_central']
+g16 = ['p_consonant','n_consonant']
+g17 = ['p_sonorant','n_sonorant']
+g18 = ['p_long','n_long']
+g19 = ['p_short','n_short']
+g20 = ['p_vowel','n_vowel']
+g21 = ['p_semivowel','n_semivowel']
+g22 = ['p_fricative','n_fricative']
+g23 = ['p_nasal','n_nasal']
+g24 = ['p_stop','n_stop']
+g25 = ['p_approximant','n_approximant']
+g26 = ['p_affricate','n_affricate']
+g27 = ['p_liquid','n_liquid']
+g28 = ['p_continuant','n_continuant']
+g29 = ['p_monophthong','n_monophthong']
+g30 = ['p_diphthong','n_diphthong']
+g31 = ['p_round','n_round']
+g32 = ['p_voiced','n_voiced']
+g33 = ['p_bilabial','n_bilabial']
+g34 = ['p_coronal','n_coronal']
+g35 = ['p_dorsal','n_dorsal']
+groups = [g1,g2,g3,g4,g5,g6,g7,g8,g9,g10,g11,g12,g13,g14,g15,g16,g17,g18,g19,g20,g21,g22,g23,g24,g25,g26,g27,g28,g29,g30,g31,g32,g33,g34,g35]
+
+
+model_dir = 'model/'
+processor = Wav2Vec2Processor.from_pretrained(model_dir)
+model = Wav2Vec2ForCTC.from_pretrained(model_dir)
+tokenizer_phoneme = Wav2Vec2CTCTokenizer(join(model_dir,"phoneme_vocab.json"), pad_token="<pad>", word_delimiter_token="")
+
+phoneme_list = list(tokenizer_phoneme.get_vocab().keys())
+p_att = pd.read_csv(join(model_dir,"phonological_attributes_v12.csv"),index_col=0)
+
+mappers = []
+for g in groups:
+    p2att = {}
+    for att in g:
+        att_phs = p_att[p_att[att]==1].index
+        for ph in att_phs:
+            p2att[ph] = att
+    mappers.append(p2att)
+
+p2att = torch.zeros((tokenizer_phoneme.vocab_size, processor.tokenizer.vocab_size)).type(torch.FloatTensor)
+
+for p in phoneme_list:
+    for mapper in mappers:
+        if p == processor.tokenizer.pad_token:
+            p2att[tokenizer_phoneme.convert_tokens_to_ids(p),processor.tokenizer.pad_token_id] = 1
+        else:
+            p2att[tokenizer_phoneme.convert_tokens_to_ids(p), processor.tokenizer.convert_tokens_to_ids(mapper[p])] = 1
+
+
+group_ids = [sorted(processor.tokenizer.convert_tokens_to_ids(group)) for group in groups]
+group_ids = [dict([(x[0]+1,x[1]) for x in list(enumerate(g))]) for g in group_ids] #This is the inversion of the one used in training as here we need to map prediction back to original tokens
+def masked_log_softmax(vector: torch.Tensor, mask: torch.Tensor, dim: int = -1) -> torch.Tensor:
+    if mask is not None:
+        mask = mask.float()
+        while mask.dim() < vector.dim():
+            mask = mask.unsqueeze(1)
+        # vector + mask.log() is an easy way to zero out masked elements in logspace, but it
+        # results in nans when the whole vector is masked.  We need a very small value instead of a
+        # zero in the mask for these cases.  log(1 + 1e-45) is still basically 0, so we can safely
+        # just add 1e-45 before calling mask.log().  We use 1e-45 because 1e-46 is so small it
+        # becomes 0 - this is just the smallest value we can actually use.
+        vector = vector + (mask + 1e-45).log()
+    return torch.nn.functional.log_softmax(vector, dim=dim)
+
+def getPhonemes(logits):
+    ngroups = len(group_ids)
+    log_props_all_masked = []
+    for i in range(ngroups):
+        mask = torch.zeros(logits.size()[2], dtype = torch.bool)
+        mask[0] = True
+        mask[list(group_ids[i].values())] = True
+        mask.unsqueeze_(0).unsqueeze_(0)
+        log_probs = masked_log_softmax(vector=logits, mask=mask, dim=-1).masked_fill(~mask,0)
+        log_props_all_masked.append(log_probs)
+    log_probs_cat = torch.stack(log_props_all_masked, dim=0).sum(dim=0)
+    log_probs_phoneme = torch.matmul(p2att,log_probs_cat.transpose(1,2)).transpose(1,2).type(torch.FloatTensor)
+    pred_ids = torch.argmax(log_probs_phoneme,dim=-1)
+    pred = tokenizer_phoneme.batch_decode(pred_ids,spaces_between_special_tokens=True)[0]
+    return pred
+
+def getAtt(logits,i):
+    mask = torch.zeros(logits.size()[2], dtype = torch.bool)
+    mask[0] = True
+    mask[list(group_ids[i].values())] = True
+    logits_g = logits[:,:,mask]
+    pred_ids = torch.argmax(logits_g,dim=-1)
+    pred_ids = pred_ids.cpu().apply_(lambda x: group_ids[i].get(x,x))
+    pred = processor.batch_decode(pred_ids,spaces_between_special_tokens=True)[0]
+    return pred.replace('p_','+').replace('n_','-')
+
+def recognizeAudio(audio, mic_audioFilePath, att):
+    i = Attributes[att]
+    audio = mic_audioFilePath if mic_audioFilePath else audio
+    y, sr = librosa.load(audio, sr=16000)
+    input_values = processor(audio=y, sampling_rate=sr, return_tensors="pt").input_values
+    with torch.no_grad():
+        logits = model(input_values).logits
+    return getPhonemes(logits), getAtt(logits,i)
+
+
+gui = gr.Interface(fn=recognizeAudio, inputs=[Audio(label="Upload Audio File", type="filepath"),Audio(source="microphone", type="filepath", label="Record from microphone"), 
+                                               Dropdown(choices=Attributes.keys(),type="value",label="Select Attribute")],
+                    outputs=[Textbox(label="ARPA Phoneme"),Textbox(label="Attribute (+/-)")])
+    
+gui.launch()