Update app.py

app.py

@@ -4,19 +4,21 @@ import torch
 import torchaudio
 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,
     TrainingArguments,
     Trainer,
-    DataCollatorForSeq2Seq,  # ✅ Fix: Use correct data collator
+    DataCollatorForSeq2Seq,
 )
 
 # ================================
 # 1️⃣ Authenticate with Hugging Face Hub (Securely)
 # ================================
-HF_TOKEN = os.getenv("hf_token")  # Ensure it's set in Hugging Face Spaces Secrets
+HF_TOKEN = os.getenv("hf_token")  
 
 if HF_TOKEN is None:
     raise ValueError("❌ Hugging Face API token not found. Please set it in Secrets.")
@@ -30,18 +32,16 @@ MODEL_NAME = "AqeelShafy7/AudioSangraha-Audio_to_Text"
 processor = AutoProcessor.from_pretrained(MODEL_NAME)
 model = AutoModelForSpeechSeq2Seq.from_pretrained(MODEL_NAME)
 
-# Move model to GPU if available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 print(f"✅ Model loaded on {device}")
 
 # ================================
-# 3️⃣ Load Dataset (Recursively from Extracted Path)
+# 3️⃣ Load and Prepare Dataset
 # ================================
 DATASET_TAR_PATH = "dev-clean.tar.gz"
 EXTRACT_PATH = "./librispeech_dev_clean"
 
-# 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:
@@ -50,58 +50,42 @@ if not os.path.exists(EXTRACT_PATH):
 else:
     print("✅ Dataset already extracted.")
 
-# Base directory where audio files are stored
-AUDIO_FOLDER = os.path.join(EXTRACT_PATH, "LibriSpeech", "dev-clean")
+# Load dataset with transcripts
+dataset = load_dataset("librispeech_asr", "clean", split="train")
 
-# Recursively find all `.flac` files inside the dataset directory
-def find_audio_files(base_folder):
-    """Recursively search for all .flac files in subdirectories."""
-    audio_files = []
-    for root, _, files in os.walk(base_folder):
-        for file in files:
-            if file.endswith(".flac"):
-                audio_files.append(os.path.join(root, file))
-    return audio_files
+# Ensure dataset has transcripts
+if "text" not in dataset.column_names:
+    raise ValueError("❌ Dataset is missing transcription text!")
 
-# Get all audio files
-audio_files = find_audio_files(AUDIO_FOLDER)
-
-if not audio_files:
-    raise FileNotFoundError(f"❌ No .flac files found in {AUDIO_FOLDER}. Check dataset structure!")
-
-print(f"✅ Found {len(audio_files)} audio files in dataset!")
-
-# ================================
-# 4️⃣ Preprocess Dataset (Fixed input_features)
-# ================================
-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
+# Preprocessing Function
+def preprocess_data(batch):
+    # Process audio
+    waveform, sample_rate = torchaudio.load(batch["file"])
     waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(waveform)
-
-    # Convert to model input format
-    input_features = processor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_features[0]
-
-    return input_features
-
-# Manually create dataset structure
-dataset = [{"input_features": load_and_process_audio(f), "labels": []} for f in audio_files[:100]]
-
-# Split dataset into train and eval
+    
+    batch["input_features"] = processor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_features[0]
+    
+    # Tokenize transcript text
+    batch["labels"] = processor.tokenizer(batch["text"], padding="max_length", truncation=True, return_tensors="pt").input_ids[0]
+    
+    return batch
+
+# Apply preprocessing
+dataset = dataset.map(preprocess_data, remove_columns=["file", "audio", "text"])
+
+# Split into train & eval
 train_size = int(0.8 * len(dataset))
-train_dataset = dataset[:train_size]
-eval_dataset = dataset[train_size:]
+train_dataset = dataset.select(range(train_size))
+eval_dataset = dataset.select(range(train_size, len(dataset)))
 
-print(f"✅ Dataset Loaded! Training: {len(train_dataset)}, Evaluation: {len(eval_dataset)}")
+print(f"✅ Dataset Prepared! Training: {len(train_dataset)}, Evaluation: {len(eval_dataset)}")
 
 # ================================
-# 5️⃣ Training Arguments & Trainer
+# 4️⃣ Training Arguments & Trainer
 # ================================
 training_args = TrainingArguments(
     output_dir="./asr_model_finetuned",
-    eval_strategy="epoch",  # Fixed deprecated evaluation_strategy
+    evaluation_strategy="epoch",
     save_strategy="epoch",
     learning_rate=5e-5,
     per_device_train_batch_size=8,
@@ -111,15 +95,13 @@ training_args = TrainingArguments(
     logging_dir="./logs",
     logging_steps=500,
     save_total_limit=2,
-    push_to_hub=True,  # Fix: Properly authenticate Hugging Face Hub
-    hub_model_id="tahirsher/ASR_Model_for_Transcription_into_Text",  # Replace with your Hugging Face repo
+    push_to_hub=True,
+    hub_model_id="tahirsher/ASR_Model_for_Transcription_into_Text",
     hub_token=HF_TOKEN,
 )
 
-# ✅ FIX: Use correct Data Collator
 data_collator = DataCollatorForSeq2Seq(tokenizer=processor.tokenizer, model=model, return_tensors="pt")
 
-# Define Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -129,45 +111,54 @@ trainer = Trainer(
 )
 
 # ================================
-# 6️⃣ Fine-Tuning Execution
+# 5️⃣ Fine-Tuning Execution & Training Stats
 # ================================
 if st.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
+    train_loss = trainer.state.log_history
+    losses = [entry['loss'] for entry in train_loss if 'loss' in entry]
+    
+    plt.figure(figsize=(8, 5))
+    plt.plot(range(len(losses)), losses, label="Training Loss", color="blue")
+    plt.xlabel("Steps")
+    plt.ylabel("Loss")
+    plt.title("Training Loss Over Time")
+    plt.legend()
+    st.pyplot(plt)
+
 # ================================
-# 7️⃣ Streamlit ASR Web App
+# 6️⃣ Streamlit ASR Web App (Proper Decoding)
 # ================================
-st.title("🎙️ Speech-to-Text ASR with Fine-Tuning 🎶")
+st.title("🎙️ Speech-to-Text ASR Model with Fine-Tuning 🎶")
 
-# Upload audio file
 audio_file = st.file_uploader("Upload an audio file", type=["wav", "mp3", "flac"])
 
 if audio_file:
-    # Save uploaded file temporarily
     audio_path = "temp_audio.wav"
     with open(audio_path, "wb") as f:
         f.write(audio_file.read())
 
-    # Load and process audio
     waveform, sample_rate = torchaudio.load(audio_path)
     waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(waveform)
 
-    # Convert audio to model input
     input_features = processor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_features
 
-    # ✅ FIX: Ensure input tensor is correctly formatted
-    input_tensor = input_features.to(device)  # Move to GPU/CPU
-
-    # ✅ FIX: Provide decoder_input_ids
-    decoder_input_ids = torch.tensor([[model.config.decoder_start_token_id]]).to(device)
+    input_tensor = input_features.to(device)
 
-    # Perform ASR inference
+    # ✅ FIX: Use `generate()` for Proper Transcription
     with torch.no_grad():
-        logits = model(input_tensor, decoder_input_ids=decoder_input_ids).logits
-        predicted_ids = torch.argmax(logits, dim=-1)
-        transcription = processor.batch_decode(predicted_ids)[0]
+        generated_ids = model.generate(
+            input_tensor,
+            max_length=500,
+            num_beams=5,
+            do_sample=True,
+            top_k=50
+        )
+        transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
     # Display transcription
     st.success("📄 Transcription:")