Update app.py

app.py

@@ -21,6 +21,9 @@ import torch
 import librosa
 import torchaudio
 import numpy as np
+import torch
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, Wav2Vec2FeatureExtractor, Wav2Vec2CTCTokenizer
+import librosa
 
 
 url = "https://huggingface.co/MahtaFetrat/tempmodel/resolve/main/checkpoint-15-1200.zip"
@@ -34,21 +37,47 @@ output_dir = "extracted_model"
 subprocess.run(["unzip", zip_file, "-d", output_dir], check=True)
 
 
-from transformers import Wav2Vec2CTCTokenizer
+# Function for inference from an audio file path
+def infer_from_audio_file(audio_file_path, model, processor, device="cpu"):
+    # Load audio file
+    audio, sampling_rate = librosa.load(audio_file_path, sr=16000)
 
-tokenizer = Wav2Vec2CTCTokenizer("vocab.json", unk_token="<unk>", pad_token="<pad>", word_delimiter_token="|")
+    # Process the audio using the feature extractor from the processor
+    inputs = processor(audio, sampling_rate=sampling_rate).input_values[0]
+    input_features = [{"input_values": inputs}]
 
+    batch = processor.pad(
+            input_features,
+            padding=True,
+            max_length=None,
+            pad_to_multiple_of=None,
+            return_tensors="pt",
+        )
 
-from transformers import Wav2Vec2FeatureExtractor
+    # Move inputs to the correct device
+    input_values = batch.input_values.to(device)
 
-feature_extractor = Wav2Vec2FeatureExtractor(feature_size=1, sampling_rate=16000, padding_value=0.0, do_normalize=True, return_attention_mask=True)
+    # Ensure the model is in evaluation mode
+    model.eval()
 
+    with torch.no_grad():
+        # Make predictions
+        outputs = model(input_values)
+        logits = outputs.logits
 
-from transformers import Wav2Vec2Processor
+        # Decode the predictions
+        pred_ids = torch.argmax(logits, dim=-1)
+        pred_str = processor.batch_decode(pred_ids.cpu().numpy())
 
-tuned_wav2vec_processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)
-tuned_wav2vec_model = Wav2Vec2ForCTC.from_pretrained("extracted_model/checkpoint-15-1200")
+    return pred_str[0]  # Return the decoded transcription of the audio
+
+
+tokenizer = Wav2Vec2CTCTokenizer("./vocab.json", unk_token="<unk>", pad_token="<pad>", word_delimiter_token="|")
+feature_extractor = Wav2Vec2FeatureExtractor(feature_size=1, sampling_rate=16000, padding_value=0.0, do_normalize=True, return_attention_mask=True)
+processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)
 
+latest_checkpoint = "extracted_model/checkpoint-15-1200"
+model = Wav2Vec2ForCTC.from_pretrained(latest_checkpoint)
 
 def tuned_wav2vec_speech_file_to_array_fn(path):
     speech_array, sampling_rate = torchaudio.load(path)
@@ -88,34 +117,8 @@ def preprocess_audio(audio_path):
 
 
 def speech_to_text(audio_path):
-    # waveform = preprocess_audio(audio_path)
-    
-    # input_values = processor(waveform.squeeze(), return_tensors="pt", sampling_rate=16000).input_values
-    # with torch.no_grad():
-    #     logits = model(input_values).logits
-    # predicted_ids = torch.argmax(logits, dim=-1)
-    # transcription = processor.batch_decode(predicted_ids)[0]
-    # return transcription
-
-    speech = tuned_wav2vec_speech_file_to_array_fn(audio_path)
-
-    features = tuned_wav2vec_processor(
-        speech,
-        sampling_rate=tuned_wav2vec_processor.feature_extractor.sampling_rate,
-        return_tensors="pt",
-        padding=True
-    )
-
-    input_values = features.input_values
-    attention_mask = features.attention_mask
-
-    with torch.no_grad():
-        logits = tuned_wav2vec_model(input_values, attention_mask=attention_mask).logits
-
-    pred_ids = torch.argmax(logits, dim=-1)
-
-    predicted = tuned_wav2vec_processor.batch_decode(pred_ids)
-    return predicted[0]
+    predicted = infer_from_audio_file(audio_path, model, processor)
+    return predicted
 
 
 def video_to_text(video_path):