Update app.py

app.py

@@ -1,77 +1,165 @@
+import gradio as gr
 import torch
-import pickle
-import streamlit as st
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+import logging
+import traceback
+import sys
+from audio_processing import AudioProcessor
 
-from io import BytesIO
-from config import SAMPLING_RATE
-from model_utils import load_models
-from audio_processing import detect_language, process_long_audio, load_and_resample_audio
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+logger = logging.getLogger(__name__)
 
-# Clear GPU cache
-torch.cuda.empty_cache()
+def load_qa_model():
+    """Load question-answering model"""
+    try:
+        qa_pipeline = pipeline(
+            "text-generation",
+            model="meta-llama/Meta-Llama-3-8B-Instruct",
+            model_kwargs={"torch_dtype": torch.bfloat16},
+            device_map="auto",
+        )
+        return qa_pipeline
+    except Exception as e:
+        logger.error(f"Failed to load Q&A model: {str(e)}")
+        return None
 
-# Load models at startup
-load_models()
+def load_summarization_model():
+    """Load summarization model"""
+    try:
+        summarizer = pipeline(
+            "summarization", 
+            model="sshleifer/distilbart-cnn-12-6",
+            device=0 if torch.cuda.is_available() else -1
+        )
+        return summarizer
+    except Exception as e:
+        logger.error(f"Failed to load summarization model: {str(e)}")
+        return None
 
-# Title of the app
-st.title("Audio Player with Live Transcription and Translation")
+def process_audio(audio_file, translate=False):
+    """Process audio file"""
+    try:
+        processor = AudioProcessor()
+        language_segments, final_segments = processor.process_audio(audio_file, translate)
+        
+        # Format output
+        transcription = ""
+        full_text = ""
+        
+        # Add language detection information
+        for segment in language_segments:
+            transcription += f"Language: {segment['language']}\n"
+            transcription += f"Time: {segment['start']:.2f}s - {segment['end']:.2f}s\n\n"
+        
+        # Add transcription/translation information
+        transcription += "Transcription with language detection:\n\n"
+        for segment in final_segments:
+            transcription += f"[{segment['start']:.2f}s - {segment['end']:.2f}s] ({segment['language']}):\n"
+            transcription += f"Original: {segment['text']}\n"
+            if translate and 'translated' in segment:
+                transcription += f"Translated: {segment['translated']}\n"
+                full_text += segment['translated'] + " "
+            else:
+                full_text += segment['text'] + " "
+            transcription += "\n"
+        
+        return transcription, full_text
+        
+    except Exception as e:
+        logger.error(f"Audio processing failed: {str(e)}")
+        raise gr.Error(f"Processing failed: {str(e)}")
 
-# Sidebar for file uploader and submit button
-st.sidebar.header("Upload Audio Files")
-uploaded_files = st.sidebar.file_uploader("Choose audio files", type=["mp3", "wav"], accept_multiple_files=True)
-submit_button = st.sidebar.button("Submit")
+def summarize_text(text):
+    """Summarize text"""
+    try:
+        summarizer = load_summarization_model()
+        if summarizer is None:
+            return "Summarization model could not be loaded."
+        
+        summary = summarizer(text, max_length=150, min_length=50, do_sample=False)[0]['summary_text']
+        return summary
+    except Exception as e:
+        logger.error(f"Summarization failed: {str(e)}")
+        return "Error occurred during summarization."
 
-# Session state to hold data
-if 'audio_files' not in st.session_state:
-    st.session_state.audio_files = []
-    st.session_state.transcriptions = {}
-    st.session_state.translations = {}
-    st.session_state.detected_languages = []
-    st.session_state.waveforms = []
+def answer_question(context, question):
+    """Answer questions about the text"""
+    try:
+        qa_pipeline = load_qa_model()
+        if qa_pipeline is None:
+            return "Q&A model could not be loaded."
+        
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant who can answer questions based on the given context."},
+            {"role": "user", "content": f"Context: {context}\n\nQuestion: {question}"}
+        ]
+        
+        response = qa_pipeline(messages, max_new_tokens=256)[0]['generated_text']
+        return response
+    except Exception as e:
+        logger.error(f"Q&A failed: {str(e)}")
+        return f"Error occurred during Q&A process: {str(e)}"
 
-# Process uploaded files
-if submit_button and uploaded_files is not None:
-    st.session_state.audio_files = uploaded_files
-    st.session_state.detected_languages = []
-    st.session_state.waveforms = []
+# Create Gradio interface
+with gr.Blocks() as iface:
+    gr.Markdown("# Multilingual Speech Processing with MMS")
+    
+    with gr.Row():
+        with gr.Column():
+            audio_input = gr.Audio(type="filepath")
+            translate_checkbox = gr.Checkbox(label="Enable Translation")
+            process_button = gr.Button("Process Audio")
+        
+        with gr.Column():
+            transcription_output = gr.Textbox(label="Transcription/Translation", lines=10)
+            full_text_output = gr.Textbox(label="Full Text", lines=5)
+    
+    with gr.Row():
+        with gr.Column():
+            summarize_button = gr.Button("Summarize")
+            summary_output = gr.Textbox(label="Summary", lines=3)
+            
+        with gr.Column():
+            question_input = gr.Textbox(label="Ask a question about the transcription")
+            answer_button = gr.Button("Get Answer")
+            answer_output = gr.Textbox(label="Answer", lines=3)
+    
+    # Set up event handlers
+    process_button.click(
+        process_audio,
+        inputs=[audio_input, translate_checkbox],
+        outputs=[transcription_output, full_text_output]
+    )
+    
+    summarize_button.click(
+        summarize_text,
+        inputs=[full_text_output],
+        outputs=[summary_output]
+    )
+    
+    answer_button.click(
+        answer_question,
+        inputs=[full_text_output, question_input],
+        outputs=[answer_output]
+    )
+    
+    # Add system information
+    gr.Markdown(f"""
+    ## System Information
+    - Device: {"CUDA" if torch.cuda.is_available() else "CPU"}
+    - CUDA Available: {"Yes" if torch.cuda.is_available() else "No"}
+    
+    ## Features
+    - Automatic language detection
+    - High-quality transcription using MMS
+    - Optional translation to English
+    - Text summarization
+    - Question answering
+    """)
 
-    for uploaded_file in uploaded_files:
-        waveform = load_and_resample_audio(BytesIO(uploaded_file.read()))
-        st.session_state.waveforms.append(waveform)
-        detected_language = detect_language(waveform)
-        st.session_state.detected_languages.append(detected_language)
-
-# Display uploaded files and options
-if 'audio_files' in st.session_state and st.session_state.audio_files:
-    for i, uploaded_file in enumerate(st.session_state.audio_files):
-        st.write(f"**File name**: {uploaded_file.name}")
-        st.audio(uploaded_file, format=uploaded_file.type)
-        st.write(f"**Detected Language**: {st.session_state.detected_languages[i]}")
-
-        col1, col2 = st.columns(2)
-
-        with col1:
-            if st.button(f"Transcribe {uploaded_file.name}"):
-                with st.spinner("Transcribing..."):
-                    transcription = process_long_audio(st.session_state.waveforms[i], SAMPLING_RATE)
-                    st.session_state.transcriptions[i] = transcription
-
-            if st.session_state.transcriptions.get(i):
-                st.write("**Transcription**:")
-                st.text_area("", st.session_state.transcriptions[i], height=200, key=f"transcription_{i}")
-                st.markdown(f'<div style="text-align: right;"><a href="data:text/plain;charset=UTF-8,{st.session_state.transcriptions[i]}" download="transcription_{uploaded_file.name}.txt">Download Transcription</a></div>', unsafe_allow_html=True)
-
-        with col2:
-            if st.button(f"Translate {uploaded_file.name}"):
-                with st.spinner("Translating..."):
-                    with open('languages.pkl', 'rb') as f:
-                        lang_dict = pickle.load(f)
-                    detected_language_name = lang_dict[st.session_state.detected_languages[i]]
-
-                    translation = process_long_audio(st.session_state.waveforms[i], SAMPLING_RATE, task="translate",
-                                                     language=detected_language_name)
-                    st.session_state.translations[i] = translation
-            if st.session_state.translations.get(i):
-                st.write("**Translation**:")
-                st.text_area("", st.session_state.translations[i], height=200, key=f"translation_{i}")
-                st.markdown(f'<div style="text-align: right;"><a href="data:text/plain;charset=UTF-8,{st.session_state.translations[i]}" download="translation_{uploaded_file.name}.txt">Download Translation</a></div>', unsafe_allow_html=True)
+if __name__ == "__main__":
+    iface.launch()