pretrained model

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import whisper
 import librosa
 import torch
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Tokenizer
+from transformers import Wav2Vec2Processor, Wav2Vec2Tokenizer
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -14,18 +14,28 @@ def audio_to_text(audio):
     result = model.transcribe(audio)
 
     return result["text"]
-
-    # audio, rate = librosa.load(audio, sr = 16000)
-
     # tokenizer = Wav2Vec2Tokenizer.from_pretrained("facebook/wav2vec2-base-960h")
-    # model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h").to(device)
-
-    # input_values = tokenizer(audio, return_tensors="pt").input_values.to(device)
-    # logits = model(input_values).logits
 
-    # prediction = torch.argmax(logits, dim=-1)
-    # transcription = tokenizer.batch_decode(prediction)[0]
-    return transcription
+    # logits = preprocess(audio)
+
+    # 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")
+
+    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,