Update app.py

app.py

@@ -1,53 +1,61 @@
-from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
+import os
+import torch
+import torchaudio
+import tarfile
+import numpy as np
+import streamlit as st
+from datasets import load_dataset
+from transformers import (
+    AutoProcessor,
+    AutoModelForSpeechSeq2Seq,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForSeq2Seq,
+)
 
-# Load Processor & Model
-processor = AutoProcessor.from_pretrained("AqeelShafy7/AudioSangraha-Audio_to_Text")
-model = AutoModelForSpeechSeq2Seq.from_pretrained("AqeelShafy7/AudioSangraha-Audio_to_Text")
+# ================================
+# 1️⃣ Load Model & Processor
+# ================================
+MODEL_NAME = "AqeelShafy7/AudioSangraha-Audio_to_Text"
+
+# Load ASR model and processor
+processor = AutoProcessor.from_pretrained(MODEL_NAME)
+model = AutoModelForSpeechSeq2Seq.from_pretrained(MODEL_NAME)
 
 # Move model to GPU if available
-import torch
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
-print(f"Model loaded on {device}")
-
-from datasets import load_dataset
-import torchaudio
-import torch
-
-# Fix: Add trust_remote_code=True
-import fsspec
-import os
-import tarfile
+print(f"✅ Model loaded on {device}")
 
-# Define paths
-dataset_tar_path = "dev-clean.tar.gz"  # Path in your repo
-extract_path = "./librispeech_dev_clean"  # Extracted folder
+# ================================
+# 2️⃣ Load Dataset (LibriSpeech)
+# ================================
 
-# Check if dataset is already extracted, if not, extract it
-if not os.path.exists(extract_path):
-    print("Extracting dataset...")
-    with tarfile.open(dataset_tar_path, "r:gz") as tar:
-        tar.extractall(extract_path)
-    print("Extraction complete.")
-else:
-    print("Dataset already extracted.")
-
-from datasets import load_dataset
+DATASET_TAR_PATH = "dev-clean.tar.gz"  # The uploaded dataset in Hugging Face space
+EXTRACT_PATH = "./librispeech_dev_clean"  # Extracted folder
 
-# Load extracted dataset
-dataset = load_dataset("librispeech_asr", data_dir=extract_path, split="train", trust_remote_code=True)
+# Extract dataset if not already extracted
+if not os.path.exists(EXTRACT_PATH):
+    print("🔄 Extracting dataset...")
+    with tarfile.open(DATASET_TAR_PATH, "r:gz") as tar:
+        tar.extractall(EXTRACT_PATH)
+    print("✅ Extraction complete.")
 
-print("Dataset loaded successfully!")
+# Load dataset from extracted path
+dataset = load_dataset("librispeech_asr", data_dir=EXTRACT_PATH, split="train", trust_remote_code=True)
+print(f"✅ Dataset Loaded! {dataset}")
 
-# Function to load & resample audio
+# ================================
+# 3️⃣ Preprocess Dataset
+# ================================
 def preprocess_audio(batch):
     audio = batch["audio"]
     waveform, sample_rate = torchaudio.load(audio["path"])
-    
-    # Resample to 16kHz (ASR models usually require this)
+
+    # Resample to 16kHz
     waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(waveform)
-    
-    # Convert to correct format
+
+    # Convert to model input format
     batch["input_values"] = processor(waveform.squeeze().numpy(), sampling_rate=16000).input_values[0]
     batch["labels"] = processor.tokenizer(batch["text"]).input_ids
     return batch
@@ -55,9 +63,9 @@ def preprocess_audio(batch):
 # Apply preprocessing
 dataset = dataset.map(preprocess_audio, remove_columns=["audio"])
 
-from transformers import TrainingArguments, Trainer, DataCollatorForSeq2Seq
-
-# Define Training Arguments
+# ================================
+# 4️⃣ Training Arguments & Trainer
+# ================================
 training_args = TrainingArguments(
     output_dir="./asr_model_finetuned",
     evaluation_strategy="epoch",
@@ -70,10 +78,14 @@ training_args = TrainingArguments(
     logging_dir="./logs",
     logging_steps=500,
     save_total_limit=2,
-    push_to_hub=True,  # Enable uploading to Hugging Face Hub
+    push_to_hub=True,
+    metric_for_best_model="wer",
+    greater_is_better=False,
+    save_on_each_node=True,  # Improves stability during multi-GPU training
+    load_best_model_at_end=True,  # Saves best model
 )
 
-# Define Data Collator
+# Data collator (for dynamic padding)
 data_collator = DataCollatorForSeq2Seq(processor.tokenizer, model=model)
 
 # Define Trainer
@@ -81,38 +93,41 @@ trainer = Trainer(
     model=model,
     args=training_args,
     train_dataset=dataset,
-    eval_dataset=None,  # We use only training data here
+    eval_dataset=None,  # No validation dataset for now
     tokenizer=processor.feature_extractor,
     data_collator=data_collator,
 )
 
-# Start Fine-Tuning
-trainer.train()
+# ================================
+# 5️⃣ Fine-Tuning Execution
+# ================================
+if st.button("Start Fine-Tuning"):
+    with st.spinner("Fine-tuning in progress... Please wait!"):
+        trainer.train()
+    st.success("✅ Fine-Tuning Completed! Model updated.")
 
-# Deployment of Huggingface using streamlit
-import streamlit as st
-import soundfile as sf
-import numpy as np
-
-st.title("🎙️ Automatic Speech Recognition with Fine-Tuning 🎶")
+# ================================
+# 6️⃣ Streamlit ASR Web App
+# ================================
+st.title("🎙️ Speech-to-Text ASR with Fine-Tuning 🎶")
 
 # Upload audio file
 audio_file = st.file_uploader("Upload an audio file", type=["wav", "mp3", "flac"])
 
 if audio_file:
-    # Save and load audio file
-    with open("temp_audio.wav", "wb") as f:
+    # Save uploaded file temporarily
+    audio_path = "temp_audio.wav"
+    with open(audio_path, "wb") as f:
         f.write(audio_file.read())
 
-    waveform, sample_rate = torchaudio.load("temp_audio.wav")
-
-    # Resample to 16kHz
+    # Load and process audio
+    waveform, sample_rate = torchaudio.load(audio_path)
     waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(waveform)
 
-    # Convert to model input
+    # Convert audio to model input
     input_values = processor(waveform.squeeze().numpy(), sampling_rate=16000).input_values[0]
 
-    # Perform transcription
+    # Perform ASR inference
     with torch.no_grad():
         input_tensor = torch.tensor([input_values]).to(device)
         logits = model(input_tensor).logits
@@ -120,22 +135,23 @@ if audio_file:
         transcription = processor.batch_decode(predicted_ids)[0]
 
     # Display transcription
-    st.success("Transcription:")
+    st.success("📄 Transcription:")
     st.write(transcription)
 
-    # Fine-tune with user input
-    user_correction = st.text_area("Correct the transcription (if needed):")
-    
-    if st.button("Fine-Tune Model"):
+    # ================================
+    # 7️⃣ Fine-Tune Model with User Correction
+    # ================================
+    user_correction = st.text_area("🔧 Correct the transcription (if needed):", transcription)
+
+    if st.button("Fine-Tune with Correction"):
         if user_correction:
-            # Convert correction to training format
             corrected_input = processor.tokenizer(user_correction).input_ids
 
-            # Update dataset dynamically (simple approach)
+            # Dynamically add new example to dataset
             dataset = dataset.add_item({"input_values": input_values, "labels": corrected_input})
 
-            # Retrain for one step
+            # Perform quick re-training (1 epoch)
+            trainer.args.num_train_epochs = 1
             trainer.train()
 
-            st.success("Model fine-tuned successfully! Try another audio file.")
-
+            st.success("✅ Model fine-tuned with new correction! Try another audio file.")