finetuned model

app.py

@@ -1,13 +1,38 @@
+import os
 import gradio as gr
 import whisper
+import librosa
+import torch
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTCTokenizer
 
-def audio_to_text(audio):
-    model = whisper.load_model("base")
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    audio = whisper.load_audio(audio)
-    result = model.transcribe(audio)
+def audio_to_text(logits):
+    # model = whisper.load_model("base")
 
-    return result["text"]
+    # audio = whisper.load_audio(audio)
+    # result = model.transcribe(audio)
+
+    # return result["text"]
+    tokenizer = Wav2Vec2ForCTCTokenizer("model_save/wav2vec2_osr_version_1_vocab/vocab.json", unk_token="<unk>", pad_token="<pad>", word_delimiter_token="|")
+
+    predicted_ids = torch.argmax(logits, dim=-1)
+    transcriptions = tokenizer.decode(predicted_ids[0])
+    return transcriptions
+
+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")
+
+    processor = Wav2Vec2Processor.from_pretrained(pipeline_path)
+    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,