Update app.py

app.py

@@ -1,29 +1,38 @@
-import gradio as gr
+import streamlit as st
 import requests
 import Levenshtein
-import numpy as np
-from transformers import pipeline
+from io import BytesIO
+import librosa
+import torch
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+from audio_recorder_streamlit import audio_recorder
 
-# Function to securely load the Hugging Face API token
-def load_hf_token():
-    # Replace this with your actual Hugging Face API token
-    return "your_huggingface_api_token"
+@st.cache_resource
+def load_model():
+    MODEL_ID = "jonatasgrosman/wav2vec2-large-xlsr-53-arabic"
+    processor = Wav2Vec2Processor.from_pretrained(MODEL_ID)
+    model = Wav2Vec2ForCTC.from_pretrained(MODEL_ID)
+    return processor, model
 
-# Function to query the Hugging Face Inference API
-def transcribe_audio_hf(audio):
+processor, model = load_model()
+
+def transcribe_audio_hf(audio_bytes):
     """
-    Transcribes speech from an audio file using the Hugging Face Inference API.
+    Transcribes speech from an audio file using a pretrained Wav2Vec2 model.
     Args:
-        audio (numpy.array): Audio data as a numpy array.
+        audio_bytes (bytes): Audio data in bytes.
     Returns:
         str: The transcription of the speech in the audio file.
     """
-    API_URL = "https://api-inference.huggingface.co/models/jonatasgrosman/wav2vec2-large-xlsr-53-arabic"
-    headers = {"Authorization": f"Bearer {load_hf_token()}"}
-    response = requests.post(API_URL, headers=headers, data=audio.tobytes())
-    return response.json().get("text", "").strip()
+    speech_array, sampling_rate = librosa.load(BytesIO(audio_bytes), sr=16000)
+    input_values = processor(speech_array, sampling_rate=sampling_rate, return_tensors="pt", padding=True).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].strip()
+    return transcription
+
 
-# Function to calculate Levenshtein similarity
 def levenshtein_similarity(transcription1, transcription2):
     """
     Calculate the Levenshtein similarity between two transcriptions.
@@ -37,99 +46,99 @@ def levenshtein_similarity(transcription1, transcription2):
     max_len = max(len(transcription1), len(transcription2))
     return 1 - distance / max_len  # Normalize to get similarity score
 
-# Function to evaluate audio similarity
-def evaluate_audio_similarity(original_audio, user_audio):
+def evaluate_audio_similarity(original_audio_bytes, user_audio_bytes):
     """
     Compares the similarity between the transcription of an original audio file and a user's audio file.
     Args:
-        original_audio (numpy.array): Original audio data.
-        user_audio (numpy.array): User's audio data.
+        original_audio_bytes (bytes): Bytes of the original audio file.
+        user_audio_bytes (bytes): Bytes of the user's audio file.
     Returns:
         tuple: Transcriptions and Levenshtein similarity score.
     """
-    transcription_original = transcribe_audio_hf(original_audio)
-    transcription_user = transcribe_audio_hf(user_audio)
-    similarity_score = levenshtein_similarity(transcription_original, transcription_user)
-    return transcription_original, transcription_user, similarity_score
-
-# Set up the Whisper ASR model for full-context and streaming ASR
-whisper_transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-base.en")
-
-# Full-context ASR function
-def full_context_asr(audio):
-    sr, y = audio
-    y = y.astype(np.float32)
-    y /= np.max(np.abs(y))
-    return whisper_transcriber({"sampling_rate": sr, "raw": y})["text"]
-
-# Streaming ASR function
-def streaming_asr(stream, new_chunk):
-    sr, y = new_chunk
-    y = y.astype(np.float32)
-    y /= np.max(np.abs(y))
-    
-    if stream is not None:
-        stream = np.concatenate([stream, y])
-    else:
-        stream = y
-        
-    return stream, whisper_transcriber({"sampling_rate": sr, "raw": stream})["text"]
-
-# Define Gradio interface for full-context ASR
-def gradio_full_context_interface(audio):
-    if audio is not None:
-        transcription = full_context_asr(audio)
-        return transcription
-    else:
-        return "Please provide an audio file."
-
-# Define Gradio interface for audio similarity checking
-def gradio_similarity_interface(original_audio, user_audio):
-    if original_audio is not None and user_audio is not None:
-        transcription_original, transcription_user, similarity_score = evaluate_audio_similarity(original_audio, user_audio)
-        
-        result = {
-            "Original Transcription": transcription_original,
-            "User Transcription": transcription_user,
-            "Levenshtein Similarity Score": similarity_score,
-        }
-        
-        if similarity_score > 0.8:  # Adjust the threshold as needed
-            result["Feedback"] = "The pronunciation is likely correct based on transcription similarity."
-        else:
-            result["Feedback"] = "The pronunciation may be incorrect based on transcription similarity."
-        
-        return result
-    else:
-        return "Please provide both original and user audio for comparison."
-
-# Create Gradio app for full-context ASR
-full_context_demo = gr.Interface(
-    fn=gradio_full_context_interface, 
-    inputs=gr.Audio(source="microphone", type="numpy"), 
-    outputs="text", 
-    title="Full-Context ASR Demo"
-)
-
-# Create Gradio app for streaming ASR
-streaming_demo = gr.Interface(
-    fn=streaming_asr,
-    inputs=["state", gr.Audio(source="microphone", type="numpy", streaming=True)],
-    outputs=["state", "text"],
-    live=True,
-    title="Streaming ASR Demo"
-)
-
-# Create Gradio app for audio similarity checking
-similarity_demo = gr.Interface(
-    fn=gradio_similarity_interface, 
-    inputs=[
-        gr.Audio(source="upload", type="numpy", label="Original Audio"),
-        gr.Audio(source="upload", type="numpy", label="User Audio")
-    ], 
-    outputs="json", 
-    title="Audio Transcription and Similarity Checker"
-)
-
-# Launch all three demos
-gr.TabbedInterface([full_context_demo, streaming_demo, similarity_demo], ["Full-Context ASR", "Streaming ASR", "Similarity Checker"]).launch()
+    transcription_original = transcribe_audio_hf(original_audio_bytes)
+    transcription_user = transcribe_audio_hf(user_audio_bytes)
+    similarity_score_levenshtein = levenshtein_similarity(transcription_original, transcription_user)
+    return transcription_original, transcription_user, similarity_score_levenshtein
+
+st.title("Audio Transcription and Similarity Checker")
+
+# Choose between upload or record
+st.sidebar.header("Input Method")
+input_method = st.sidebar.selectbox("Choose Input Method", ["Upload", "Record"])
+
+original_audio_bytes = None
+user_audio_bytes = None
+
+if input_method == "Upload":
+    # Upload original audio file
+    original_audio = st.file_uploader("Upload Original Audio", type=["wav", "mp3"])
+    # Upload user audio file
+    user_audio = st.file_uploader("Upload User Audio", type=["wav", "mp3"])
+
+    if original_audio:
+        original_audio_bytes = original_audio.read()
+        st.audio(original_audio_bytes, format="audio/wav")
+    if user_audio:
+        user_audio_bytes = user_audio.read()
+        st.audio(user_audio_bytes, format="audio/wav")
+
+    # Add a button to perform the test
+    if original_audio_bytes and user_audio_bytes:
+        if st.button("Perform Testing"):
+            with st.spinner("Performing transcription and similarity testing..."):
+                transcription_original, transcription_user, similarity_score = evaluate_audio_similarity(original_audio_bytes, user_audio_bytes)
+
+                # Display results
+                st.markdown("---")
+                st.subheader("Transcriptions and Similarity Score")
+                st.write(f"**Original Transcription:** {transcription_original}")
+                st.write(f"**User Transcription:** {transcription_user}")
+                st.write(f"**Levenshtein Similarity Score:** {similarity_score:.2f}")
+
+                if similarity_score > 0.8:  # Adjust the threshold as needed
+                    st.success("The pronunciation is likely correct based on transcription similarity.")
+                else:
+                    st.error("The pronunciation may be incorrect based on transcription similarity.")
+
+elif input_method == "Record":
+    st.write("Record or Upload Original Audio")
+    original_audio_bytes = audio_recorder(key="original_audio_recorder")
+
+    if not original_audio_bytes:
+        original_audio = st.file_uploader("Or Upload Original Audio", type=["wav", "mp3"])
+        if original_audio:
+            original_audio_bytes = original_audio.read()
+
+    if original_audio_bytes:
+        with st.spinner("Processing original audio..."):
+            st.audio(original_audio_bytes, format="audio/wav")
+
+    st.write("Record or Upload User Audio")
+    user_audio_bytes = audio_recorder(key="user_audio_recorder")
+
+    if not user_audio_bytes:
+        user_audio = st.file_uploader("Or Upload User Audio", type=["wav", "mp3"])
+        if user_audio:
+            user_audio_bytes = user_audio.read()
+
+    if user_audio_bytes:
+        with st.spinner("Processing user audio..."):
+            st.audio(user_audio_bytes, format="audio/wav")
+
+    # Add a button to perform the test
+    if original_audio_bytes and user_audio_bytes:
+        if st.button("Perform Testing"):
+            with st.spinner("Performing transcription and similarity testing..."):
+                transcription_original, transcription_user, similarity_score = evaluate_audio_similarity(original_audio_bytes, user_audio_bytes)
+
+                # Display results
+                st.markdown("---")
+                st.subheader("Transcriptions and Similarity Score")
+                st.write(f"**Original Transcription:** {transcription_original}")
+                st.write(f"**User Transcription:** {transcription_user}")
+                st.write(f"**Levenshtein Similarity Score:** {similarity_score:.2f}")
+
+                if similarity_score > 0.8:  # Adjust the threshold as needed
+                    st.success("The pronunciation is likely correct based on transcription similarity.")
+                else:
+                    st.error("The pronunciation may be incorrect based on transcription similarity.")