Update app.py

app.py

@@ -1,10 +1,45 @@
 import gradio as gr
+import torch
+import torchaudio
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+
+# Load a smaller Wav2Vec model and processor for Persian
+model_name = "facebook/wav2vec2-base"  # Smaller model
+processor = Wav2Vec2Processor.from_pretrained(model_name)
+model = Wav2Vec2ForCTC.from_pretrained(model_name)
+
+def transcribe_audio(audio):
+    # Load the audio file and resample to 16kHz
+    waveform, sample_rate = torchaudio.load(audio)
+    resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
+    waveform = resampler(waveform)
+
+    # Preprocess the audio
+    input_values = processor(waveform.squeeze().numpy(), return_tensors="pt", sampling_rate=16000).input_values
+
+    # Perform inference
+    with torch.no_grad():
+        logits = model(input_values).logits
+
+    # Decode the logits to text
+    predicted_ids = torch.argmax(logits, dim=-1)
+    transcription = processor.decode(predicted_ids[0])
+
+    return transcription
 
 with gr.Blocks(fill_height=True) as demo:
     with gr.Sidebar():
         gr.Markdown("# Inference Provider")
         gr.Markdown("This Space showcases the google/gemma-2-2b-it model, served by the nebius API. Sign in with your Hugging Face account to use this API.")
         button = gr.LoginButton("Sign in")
-    gr.load("models/google/gemma-2-2b-it", accept_token=button, provider="nebius")
     
+    with gr.Tab("Text Inference"):
+        gr.load("models/google/gemma-2-2b-it", accept_token=button, provider="nebius")
+    
+    with gr.Tab("Persian ASR"):
+        audio_input = gr.Audio(label="Upload Audio", type="filepath")
+        text_output = gr.Textbox(label="Transcription")
+        transcribe_button = gr.Button("Transcribe")
+        transcribe_button.click(transcribe_audio, inputs=audio_input, outputs=text_output)
+
 demo.launch()