Update app.py

app.py

@@ -1,40 +1,46 @@
-import streamlit as st
-import tensorflow as tf
-import numpy as np
 import librosa
-from datetime import datetime
+import numpy as np
 
-# Load models
-speech_to_text_model = tf.keras.models.load_model('speech_to_text_model.h5')
-translation_model = tf.keras.models.load_model('translation_model.h5')
+def preprocess_audio(file_path):
+    y, sr = librosa.load(file_path, sr=16000)
+    mel_spectrogram = librosa.feature.melspectrogram(y=y, sr=sr)
+    return mel_spectrogram
 
-def preprocess_audio(file):
-    audio, sr = librosa.load(file, sr=16000)
-    mfccs = librosa.feature.mfcc(y=audio, sr=sr, n_mfcc=13)
-    return np.expand_dims(mfccs, axis=0)
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Tokenizer, MarianMTModel, MarianTokenizer
 
-def translate_speech_to_text(audio_file):
-    audio_features = preprocess_audio(audio_file)
-    predicted_text = speech_to_text_model.predict(audio_features)
-    translated_text = translation_model.predict([predicted_text])
-    return translated_text
+# Load pre-trained models
+speech_to_text_model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-xlsr-53")
+speech_to_text_tokenizer = Wav2Vec2Tokenizer.from_pretrained("facebook/wav2vec2-large-xlsr-53")
 
-def is_after_six_pm():
-    current_time = datetime.now()
-    return current_time.hour >= 18
+translation_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-hi")
+translation_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-hi")
 
-def main():
-    st.title("Audio Translation App")
+def translate_audio(file_path):
+    # Preprocess the audio
+    mel_spectrogram = preprocess_audio(file_path)
     
-    uploaded_file = st.file_uploader("Choose an audio file", type="wav")
+    # Convert to text
+    audio_input = speech_to_text_tokenizer(file_path, return_tensors="pt").input_values
+    logits = speech_to_text_model(audio_input).logits
+    predicted_ids = logits.argmax(dim=-1)
+    transcription = speech_to_text_tokenizer.batch_decode(predicted_ids)[0]
     
-    if uploaded_file is not None:
-        if is_after_six_pm():
-            st.write("Processing...")
-            translated_text = translate_speech_to_text(uploaded_file)
-            st.write("Translated Text:", translated_text)
-        else:
-            st.write("Service available only after 6 PM IST.")
-
-if __name__ == "__main__":
-    main()
+    # Translate text
+    translation_input = translation_tokenizer(transcription, return_tensors="pt")
+    translated_output = translation_model.generate(**translation_input)
+    translation = translation_tokenizer.batch_decode(translated_output, skip_special_tokens=True)[0]
+    
+    return translation
+
+import datetime
+
+def should_translate():
+    now = datetime.datetime.now()
+    return now.hour >= 18
+
+def handle_translation(file_path):
+    if should_translate():
+        return translate_audio(file_path)
+    else:
+        return "Translation is only available after 6 PM IST."
+