Update app.py

app.py

@@ -6,7 +6,6 @@ import numpy as np
 import streamlit as st
 import matplotlib.pyplot as plt
 from huggingface_hub import login
-from datasets import load_dataset, DatasetDict
 from transformers import (
     AutoProcessor,
     AutoModelForSpeechSeq2Seq,
@@ -37,7 +36,7 @@ model.to(device)
 print(f"✅ Model loaded on {device}")
 
 # ================================
-# 3️⃣ Load Dataset (Recursively from Extracted Path)
+# 3️⃣ Load Dataset (From Extracted Folder)
 # ================================
 DATASET_TAR_PATH = "dev-clean.tar.gz"
 EXTRACT_PATH = "./librispeech_dev_clean"
@@ -69,34 +68,75 @@ if not audio_files:
 print(f"✅ Found {len(audio_files)} audio files in dataset!")
 
 # ================================
-# 4️⃣ Preprocess Dataset
+# 4️⃣ Load Transcripts
+# ================================
+def load_transcripts():
+    """Loads transcript text files and maps them to audio files."""
+    transcript_dict = {}
+    for root, _, files in os.walk(AUDIO_FOLDER):
+        for file in files:
+            if file.endswith(".txt"):  # Transcript files
+                with open(os.path.join(root, file), "r", encoding="utf-8") as f:
+                    for line in f:
+                        parts = line.strip().split(" ", 1)
+                        if len(parts) == 2:
+                            file_id, text = parts
+                            transcript_dict[file_id] = text
+    return transcript_dict
+
+transcripts = load_transcripts()
+if not transcripts:
+    raise FileNotFoundError("❌ No transcripts found! Check dataset structure.")
+
+print(f"✅ Loaded {len(transcripts)} transcripts.")
+
+# ================================
+# 5️⃣ Preprocess Dataset (Fixing `input_ids` issue)
 # ================================
 def load_and_process_audio(audio_path):
     """Loads and processes a single audio file into model format."""
     waveform, sample_rate = torchaudio.load(audio_path)
     waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(waveform)
+    
     input_features = processor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_features[0]
     return input_features
 
-dataset = [{"input_features": load_and_process_audio(f), "labels": []} for f in audio_files[:100]]
+dataset = []
+for audio_file in audio_files[:100]:  # Limit to 100 for faster processing
+    file_id = os.path.basename(audio_file).replace(".flac", "")
+    
+    if file_id in transcripts:
+        input_features = load_and_process_audio(audio_file)
+        labels = processor.tokenizer(transcripts[file_id], padding="max_length", truncation=True, return_tensors="pt").input_ids[0]
+
+        dataset.append({"input_features": input_features, "labels": labels})
 
 train_size = int(0.8 * len(dataset))
 train_dataset = dataset[:train_size]
 eval_dataset = dataset[train_size:]
 
-print(f"✅ Dataset Loaded! Training: {len(train_dataset)}, Evaluation: {len(eval_dataset)}")
+print(f"✅ Dataset Prepared! Training: {len(train_dataset)}, Evaluation: {len(eval_dataset)}")
+
+# ================================
+# 6️⃣ Streamlit UI: Fine-Tuning Hyperparameter Selection
+# ================================
+st.sidebar.title("🔧 Fine-Tuning Hyperparameters")
+
+num_epochs = st.sidebar.slider("Epochs", min_value=1, max_value=10, value=3)
+learning_rate = st.sidebar.select_slider("Learning Rate", options=[5e-4, 1e-4, 5e-5, 1e-5], value=5e-5)
+batch_size = st.sidebar.select_slider("Batch Size", options=[2, 4, 8, 16], value=8)
 
 # ================================
-# 4️⃣ Training Arguments & Trainer
+# 7️⃣ Training Arguments & Trainer
 # ================================
 training_args = TrainingArguments(
     output_dir="./asr_model_finetuned",
     evaluation_strategy="epoch",
     save_strategy="epoch",
-    learning_rate=5e-5,
-    per_device_train_batch_size=8,
-    per_device_eval_batch_size=8,
-    num_train_epochs=3,
+    learning_rate=learning_rate,
+    per_device_train_batch_size=batch_size,
+    per_device_eval_batch_size=batch_size,
+    num_train_epochs=num_epochs,
     weight_decay=0.01,
     logging_dir="./logs",
     logging_steps=500,
@@ -117,14 +157,14 @@ trainer = Trainer(
 )
 
 # ================================
-# 5️⃣ Fine-Tuning Execution & Training Stats
+# 8️⃣ Fine-Tuning Execution & Training Stats
 # ================================
-if st.button("Start Fine-Tuning"):
+if st.sidebar.button("🚀 Start Fine-Tuning"):
     with st.spinner("Fine-tuning in progress... Please wait!"):
         trainer.train()
     st.success("✅ Fine-Tuning Completed! Model updated.")
 
-    # Plot Training Loss
+    # ✅ Plot Training Loss
     train_loss = trainer.state.log_history
     losses = [entry['loss'] for entry in train_loss if 'loss' in entry]
     
@@ -137,7 +177,7 @@ if st.button("Start Fine-Tuning"):
     st.pyplot(plt)
 
 # ================================
-# 6️⃣ Streamlit ASR Web App (Proper Decoding)
+# 9️⃣ Streamlit ASR Web App (Proper Decoding)
 # ================================
 st.title("🎙️ Speech-to-Text ASR Model with Fine-Tuning 🎶")
 
@@ -166,6 +206,5 @@ if audio_file:
         )
         transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-    # Display transcription
     st.success("📄 Transcription:")
     st.write(transcription)