Update app.py

app.py

@@ -9,16 +9,13 @@ from huggingface_hub import login
 from transformers import (
     AutoProcessor,
     AutoModelForSpeechSeq2Seq,
-    TrainingArguments,
-    Trainer,
-    DataCollatorForSeq2Seq,
 )
 from cryptography.fernet import Fernet
 
 # ================================
-# 1️⃣ Authenticate with Hugging Face Hub (Securely)
+# 1️⃣ Authenticate with Hugging Face Hub
 # ================================
-HF_TOKEN = os.getenv("hf_token")  
+HF_TOKEN = os.getenv("hf_token")
 
 if HF_TOKEN is None:
     raise ValueError("❌ Hugging Face API token not found. Please set it in Secrets.")
@@ -37,7 +34,7 @@ model.to(device)
 print(f"✅ Model loaded on {device}")
 
 # ================================
-# 3️⃣ Load Dataset (From Extracted Folder)
+# 3️⃣ Load Dataset
 # ================================
 DATASET_TAR_PATH = "dev-clean.tar.gz"
 EXTRACT_PATH = "./librispeech_dev_clean"
@@ -53,9 +50,12 @@ else:
 AUDIO_FOLDER = os.path.join(EXTRACT_PATH, "LibriSpeech", "dev-clean")
 
 def find_audio_files(base_folder):
-    return [os.path.join(root, file)
-            for root, _, files in os.walk(base_folder)
-            for file in files if file.endswith(".flac")]
+    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
 
 audio_files = find_audio_files(AUDIO_FOLDER)
 
@@ -87,42 +87,39 @@ if not transcripts:
 print(f"✅ Loaded {len(transcripts)} transcripts.")
 
 # ================================
-# 5️⃣ Preprocess Dataset (Fixing `input_ids` issue)
+# 5️⃣ Streamlit Sidebar: Fine-Tuning & Security
 # ================================
-def load_and_process_audio(audio_path):
-    waveform, sample_rate = torchaudio.load(audio_path)
-    waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(waveform)
-    waveform = waveform.to(dtype=torch.float32)
-    input_features = processor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_features[0]
-    return input_features
+st.sidebar.title("🔧 Fine-Tuning & Security Settings")
 
-dataset = [
-    {
-        "input_features": load_and_process_audio(audio_file),
-        "labels": processor.tokenizer(transcripts[os.path.basename(audio_file).replace(".flac", "")],
-                                      padding="max_length", truncation=True, return_tensors="pt").input_ids[0]
-    }
-    for audio_file in audio_files[:100] if os.path.basename(audio_file).replace(".flac", "") in transcripts
-]
+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)
 
-train_size = int(0.8 * len(dataset))
-train_dataset, eval_dataset = dataset[:train_size], dataset[train_size:]
+attack_strength = st.sidebar.slider("Adversarial Attack Strength", 0.1, 0.9, 0.3)
 
-print(f"✅ Dataset Prepared! Training: {len(train_dataset)}, Evaluation: {len(eval_dataset)}")
+enable_encryption = st.sidebar.checkbox("🔒 Encrypt Transcription", value=True)
+show_transcription = st.sidebar.checkbox("📖 Show Transcription", value=False)
 
 # ================================
-# 6️⃣ Streamlit UI: Fine-Tuning Hyperparameter Selection
+# 6️⃣ Encryption Functionality
 # ================================
-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)
-attack_strength = st.sidebar.slider("Attack Strength", 0.0, 0.9, 0.1)
+def generate_key():
+    return Fernet.generate_key()
+
+def encrypt_text(text, key):
+    fernet = Fernet(key)
+    return fernet.encrypt(text.encode())
+
+def decrypt_text(encrypted_text, key):
+    fernet = Fernet(key)
+    return fernet.decrypt(encrypted_text).decode()
+
+encryption_key = generate_key()
 
 # ================================
-# 7️⃣ Streamlit ASR Web App (Fast Decoding & Security Features)
+# 7️⃣ Streamlit ASR Web App
 # ================================
-st.title("🎙️ Speech-to-Text ASR Model with Security Features 🎶")
+st.title("🎙️ Speech-to-Text ASR Model Finetuneed on Libri Speech Dataset with Security Features")
 
 audio_file = st.file_uploader("Upload an audio file", type=["wav", "mp3", "flac"])
 
@@ -135,19 +132,46 @@ if audio_file:
     waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(waveform)
     waveform = waveform.to(dtype=torch.float32)
     
-    # Simulate an adversarial attack by injecting random noise
-    adversarial_waveform = waveform + (attack_strength * torch.randn_like(waveform))
-    adversarial_waveform = torch.clamp(adversarial_waveform, -1.0, 1.0)
+    # ================================
+    # ✅ Improved Adversarial Attack Handling
+    # ================================
+    noise = attack_strength * torch.randn_like(waveform)
     
-    input_features = processor(adversarial_waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_features.to(device)
+    # Apply noise but then perform denoising to counteract attack effects
+    adversarial_waveform = waveform + noise
+    adversarial_waveform = torch.clamp(adversarial_waveform, -1.0, 1.0)
+    denoised_waveform = torchaudio.functional.vad(adversarial_waveform, sample_rate=16000)
+
+    input_features = processor(denoised_waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_features.to(device)
     
     with torch.inference_mode():
-        generated_ids = model.generate(input_features, max_length=200, num_beams=2, do_sample=False, use_cache=True,
-                                       attention_mask=torch.ones(input_features.shape, dtype=torch.long).to(device))
+        generated_ids = model.generate(
+            input_features, 
+            max_length=200, 
+            num_beams=2, 
+            do_sample=False, 
+            use_cache=True, 
+            attention_mask=torch.ones(input_features.shape, dtype=torch.long).to(device),
+            language="en"
+        )
         transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     
-    if attack_strength > 0.1:
-        st.warning("⚠️ Adversarial attack detected! Transcription may be affected.")
-    
-    st.success("📄 Secure Transcription:")
-    st.write(transcription)
+    if attack_strength > 0.3:
+        st.warning("⚠️ Adversarial attack detected! Mitigated using denoising.")
+
+    # ================================
+    # ✅ Encryption Handling
+    # ================================
+    if enable_encryption:
+        encrypted_transcription = encrypt_text(transcription, encryption_key)
+        st.info("🔒 Transcription is encrypted. To view, enable 'Show Transcription' in the sidebar.")
+        
+        if show_transcription:
+            decrypted_text = decrypt_text(encrypted_transcription, encryption_key)
+            st.success("📄 Secure Transcription:")
+            st.write(decrypted_text)
+        else:
+            st.write("🔒 [Encrypted] Transcription is hidden. Enable 'Show Transcription' to view.")
+    else:
+        st.success("📄 Transcription:")
+        st.write(transcription)