import gradio as gr from audio_processing import process_audio from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM import spaces import torch import logging import traceback import sys logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[ logging.StreamHandler(sys.stdout) ] ) logger = logging.getLogger(__name__) def load_translation_model() : tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M") model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M") return model, tokenizer def alternate_translation(translation_model, translation_tokenizer, inputs): # model, tokenizer = load_translation_model() tokenized_inputs = translation_tokenizer(inputs, return_tensors='pt') answer = "" # for translated_tokens = translation_model.generate(**tokenized_inputs, forced_bos_token_id=translation_tokenizer.convert_tokens_to_ids("eng_Latn"), max_length=100) return translation_tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0] def load_qa_model(): logger.info("Loading Q&A model...") try: model_id = "meta-llama/Meta-Llama-3-8B-Instruct" qa_pipeline = pipeline( "text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, device_map="auto", ) logger.info(f"Q&A model loaded successfully") return qa_pipeline except Exception as e: logger.warning(f"Failed to load Q&A model. Error: \n{str(e)}") return None def load_summarization_model(): logger.info("Loading summarization model...") try: cuda_available = torch.cuda.is_available() summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=0 if cuda_available else -1) logger.info(f"Summarization model loaded successfully on {'GPU' if cuda_available else 'CPU'}") return summarizer except Exception as e: logger.warning(f"Failed to load summarization model on GPU. Falling back to CPU. Error: \n{str(e)}") summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=-1) logger.info("Summarization model loaded successfully on CPU") return summarizer def process_with_fallback(func, *args, **kwargs): try: return func(*args, **kwargs) except Exception as e: logger.error(f"Error during processing: {str(e)}") logger.error(traceback.format_exc()) if "CUDA" in str(e) or "GPU" in str(e): logger.info("Falling back to CPU processing...") kwargs['use_gpu'] = False return func(*args, **kwargs) else: raise @spaces.GPU(duration=60) def transcribe_audio(audio_file, translate, model_size): logger.info(f"Starting transcription: translate={translate}, model_size={model_size}") try: result = process_with_fallback(process_audio, audio_file, translate=translate, model_size=model_size) # use_diarization=use_diarization logger.info("Transcription completed successfully") return result except Exception as e: logger.error(f"Transcription failed: {str(e)}") raise gr.Error(f"Transcription failed: {str(e)}") @spaces.GPU(duration=60) def summarize_text(text): logger.info("Starting text summarization") try: summarizer = load_summarization_model() summary = summarizer(text, max_length=150, min_length=50, do_sample=False)[0]['summary_text'] logger.info("Summarization completed successfully") return summary except Exception as e: logger.error(f"Summarization failed: {str(e)}") logger.error(traceback.format_exc()) return "Error occurred during summarization. Please try again." @spaces.GPU(duration=60) def process_and_summarize(audio_file, translate, model_size, do_summarize=True): logger.info(f"Starting process_and_summarize: translate={translate}, model_size={model_size}, do_summarize={do_summarize}") try: language_segments, final_segments = transcribe_audio(audio_file, translate, model_size) translation_model, translation_tokenizer = load_translation_model() # transcription = "Detected language changes:\n\n" transcription = "" for segment in language_segments: transcription += f"Language: {segment['language']}\n" transcription += f"Time: {segment['start']:.2f}s - {segment['end']:.2f}s\n\n" transcription += f"Transcription with language detection and speaker diarization (using {model_size} model):\n\n" full_text = "" for segment in final_segments: transcription += f"[{segment['start']:.2f}s - {segment['end']:.2f}s] ({segment['language']}) {segment['speaker']}:\n" transcription += f"Original: {segment['text']}\n" if translate: alt_trans=alternate_translation(translation_model, translation_tokenizer, segment['text']) transcription += f"Translated:{alt_trans}" full_text += alt_trans # transcription += f"Translated: {segment['translated']}\n" # full_text += segment['translated'] + " " else: full_text += segment['text'] + " " transcription += "\n" summary = summarize_text(full_text) if do_summarize else "" logger.info("Process and summarize completed successfully") return transcription, full_text, summary except Exception as e: logger.error(f"Process and summarize failed: {str(e)}\n") logger.error(traceback.format_exc()) raise gr.Error(f"Processing failed: {str(e)}") @spaces.GPU(duration=60) def answer_question(context, question): logger.info("Starting Q&A process") try: qa_pipeline = load_qa_model() if qa_pipeline is None: return "Error: 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}"}, ] alternate_system_message = """ You are an AI assistant designed to analyze speech transcriptions in a safe and responsible manner. Your purpose is to assist people, not to monitor or detect threats. When responding to user queries, your primary goals are: 1. To provide factual, accurate information to the best of your abilities. 2. To guide users towards appropriate resources and authorities if they are facing an emergency or urgent situation. 3. To refrain from speculating about or escalating potentially concerning situations without clear justification. 4. To avoid making judgements or taking actions that could infringe on individual privacy or civil liberties. However, if the speech suggests someone may be in immediate danger or that a crime is being planned, you should: - Identify & report - Identify any cryptic information and report it. - Avoid probing for additional details or speculating about the nature of the potential threat. - Do not provide any information that could enable or encourage harmful, illegal or unethical acts. Your role is to be a helpful, informative assistant. """ out = qa_pipeline(messages, max_new_tokens=256) logger.info(f"Raw model output: {out}") generated_text = out[0]['generated_text'] # Find the assistant's message for message in generated_text: if message['role'] == 'assistant': answer = message['content'] break else: answer = "No assistant response found in the model's output." logger.info(f"Extracted answer: {answer}") return answer except Exception as e: logger.error(f"Q&A process failed: {str(e)}") logger.error(traceback.format_exc()) return f"Error occurred during Q&A process. Please try again. Error: {str(e)}" # Main interface with gr.Blocks() as iface: gr.Markdown("# WhisperX Audio Transcription, Translation, Summarization, and Q&A (with ZeroGPU support)") audio_input = gr.Audio(type="filepath") translate_checkbox = gr.Checkbox(label="Enable Translation") summarize_checkbox = gr.Checkbox(label="Enable Summarization", interactive=False) model_dropdown = gr.Dropdown(choices=["tiny", "base", "small", "medium", "large", "large-v2", "large-v3"], label="Whisper Model Size", value="small") process_button = gr.Button("Process Audio") transcription_output = gr.Textbox(label="Transcription/Translation") full_text_output = gr.Textbox(label="Full Text") summary_output = gr.Textbox(label="Summary") question_input = gr.Textbox(label="Ask a question about the transcription") answer_button = gr.Button("Get Answer") answer_output = gr.Textbox(label="Answer") translate_alternate = gr.Button("Alternate Translation") translate_alternate_output = gr.Textbox(label="Alternate Translation") def update_summarize_checkbox(translate): return gr.Checkbox(interactive=translate) translate_checkbox.change(update_summarize_checkbox, inputs=[translate_checkbox], outputs=[summarize_checkbox]) process_button.click( process_and_summarize, inputs=[audio_input, translate_checkbox, model_dropdown, summarize_checkbox], outputs=[transcription_output, full_text_output, summary_output] ) answer_button.click( answer_question, inputs=[full_text_output, question_input], outputs=[answer_output] ) translate_alternate.click( alternate_translation, inputs=[summary_output], outputs=[translate_alternate_output] ) gr.Markdown( f""" ## System Information - Device: {"CUDA" if torch.cuda.is_available() else "CPU"} - CUDA Available: {"Yes" if torch.cuda.is_available() else "No"} ## ZeroGPU Support This application supports ZeroGPU for Hugging Face Spaces pro users. GPU-intensive tasks are automatically optimized for better performance when available. """ ) iface.launch()