Update app.py

app.py

@@ -4,7 +4,6 @@ import torch
 import torchaudio
 import numpy as np
 import streamlit as st
-from datasets import load_dataset
 from transformers import (
     AutoProcessor,
     AutoModelForSpeechSeq2Seq,
@@ -28,9 +27,9 @@ model.to(device)
 print(f"✅ Model loaded on {device}")
 
 # ================================
-# 2️⃣ Load Dataset (LibriSpeech) from Extracted Path
+# 2️⃣ Load Dataset (Manually from Extracted Path)
 # ================================
-DATASET_TAR_PATH = "dev-clean.tar.gz"  # Uploaded dataset in your Hugging Face Space
+DATASET_TAR_PATH = "dev-clean.tar.gz"  # Dataset stored in Hugging Face Space
 EXTRACT_PATH = "./librispeech_dev_clean"  # Extracted dataset folder
 
 # Extract dataset only if not already extracted
@@ -42,29 +41,29 @@ if not os.path.exists(EXTRACT_PATH):
 else:
     print("✅ Dataset already extracted.")
 
-# ✅ Load dataset from extracted folder
-dataset = load_dataset("librispeech_asr", data_dir=EXTRACT_PATH, split="train", trust_remote_code=True)
-print(f"✅ Dataset Loaded Successfully! Size: {len(dataset)}")
+# Load audio files manually
+AUDIO_FOLDER = os.path.join(EXTRACT_PATH, "LibriSpeech", "train-clean-100")  # Adjust as per structure
+audio_files = [os.path.join(AUDIO_FOLDER, f) for f in os.listdir(AUDIO_FOLDER) if f.endswith(".flac")]
 
 # ================================
-# 3️⃣ Preprocess Dataset
+# 3️⃣ Preprocess Dataset (Manually)
 # ================================
-def preprocess_audio(batch):
-    """Converts raw audio to a model-compatible format."""
-    audio = batch["audio"]
-    waveform, sample_rate = torchaudio.load(audio["path"])
+def load_and_process_audio(audio_path):
+    """Loads and processes a single audio file into model format."""
+    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 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
+    input_values = processor(waveform.squeeze().numpy(), sampling_rate=16000).input_values[0]
+
+    return input_values
+
+# Manually create dataset structure
+dataset = [{"input_values": load_and_process_audio(f), "labels": []} for f in audio_files[:100]]  # Load first 100
 
-# Apply preprocessing
-dataset = dataset.map(preprocess_audio, remove_columns=["audio"])
-print(f"✅ Dataset Preprocessed! Ready for Fine-Tuning.")
+print(f"✅ Dataset Loaded! Processed {len(dataset)} audio files.")
 
 # ================================
 # 4️⃣ Training Arguments & Trainer
@@ -82,10 +81,6 @@ training_args = TrainingArguments(
     logging_steps=500,
     save_total_limit=2,
     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
 )
 
 # Data collator (for dynamic padding)
@@ -151,7 +146,7 @@ if audio_file:
             corrected_input = processor.tokenizer(user_correction).input_ids
 
             # Dynamically add new example to dataset
-            dataset = dataset.add_item({"input_values": input_values, "labels": corrected_input})
+            dataset.append({"input_values": input_values, "labels": corrected_input})
 
             # Perform quick re-training (1 epoch)
             trainer.args.num_train_epochs = 1