Update app.py

app.py

@@ -1,32 +1,46 @@
-import librosa
+import os
 import gradio as gr
-import numpy as np
-from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
-import soundfile as sf
+import whisper
+import librosa
 import torch
+from transformers import Wav2Vec2Processor, Wav2Vec2Tokenizer
 
-# load model and tokenizer
-processor = Wav2Vec2Processor.from_pretrained("aditii09/facebook_english_asr")
-model = Wav2Vec2ForCTC.from_pretrained("aditii09/facebook_english_asr")
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def speech2text(audio):
-    sr, data = audio
+def audio_to_text(audio):
+    model = whisper.load_model("base")
 
-    # resample to 16hz
-    data_16hz = librosa.resample(data[:,0].astype(np.float32),sr,16000)
+    audio = whisper.load_audio(audio)
+    result = model.transcribe(audio)
 
-    # tokenize
-    input_values = processor([data_16hz], return_tensors="pt", padding="longest").input_values  # Batch size 1
+    return result["text"]
+    # tokenizer = Wav2Vec2Tokenizer.from_pretrained("facebook/wav2vec2-base-960h")
 
-    # retrieve logits
-    logits = model(input_values).logits
+    # logits = preprocess(audio)
 
-    # take argmax and decode
-    predicted_ids = torch.argmax(logits, dim=-1)
-    transcription = processor.batch_decode(predicted_ids)
+    # predicted_ids = torch.argmax(logits, dim=-1)
+    # transcriptions = tokenizer.decode(predicted_ids[0])
+    # return transcriptions
 
-    return transcription[0].lower()  # batch size 1
+def preprocess(audio):
+    model_save_path = "model_save"
+    model_name = "wav2vec2_osr_version_1"
+    speech, rate = librosa.load(audio, sr=16000)
+    model_path = os.path.join(model_save_path, model_name+".pt")
+    pipeline_path = os.path.join(model_save_path, model_name+"_vocab")
 
-iface = gr.Interface(speech2text, "microphone", "text")
-
-iface.launch()
+    access_token = "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"
+    processor = Wav2Vec2Processor.from_pretrained(pipeline_path, use_auth_token=access_token)
+    model = torch.load(model_path)
+    model.eval()
+    input_values = processor(speech, sampling_rate=rate, return_tensors="pt").input_values.to(device)
+    logits = model(input_values).logits
+    return logits
+
+demo = gr.Interface(
+    fn=audio_to_text,
+    inputs=gr.Audio(source="upload", type="filepath"),
+    examples=[["example.flac"]],
+    outputs="text"
+)
+demo.launch()