| import streamlit as st |
| from transformers import WhisperProcessor, WhisperForConditionalGeneration, RagTokenizer, RagRetriever, RagSequenceForGeneration |
| import torch |
| import soundfile as sf |
| import librosa |
| from moviepy.editor import VideoFileClip |
| import os |
| import tempfile |
|
|
| |
| whisper_model_name = "openai/whisper-base" |
| whisper_processor = WhisperProcessor.from_pretrained(whisper_model_name) |
| whisper_model = WhisperForConditionalGeneration.from_pretrained(whisper_model_name) |
|
|
| |
| rag_model_name = "facebook/rag-sequence-nq" |
| rag_tokenizer = RagTokenizer.from_pretrained(rag_model_name) |
| rag_retriever = RagRetriever.from_pretrained(rag_model_name, index_name="exact", use_dummy_dataset=True, trust_remote_code=True) |
| rag_model = RagSequenceForGeneration.from_pretrained(rag_model_name, retriever=rag_retriever) |
|
|
| def transcribe_audio(audio_path, language="ru"): |
| speech, rate = librosa.load(audio_path, sr=16000) |
| inputs = whisper_processor(speech, return_tensors="pt", sampling_rate=16000) |
| input_features = whisper_processor.feature_extractor(speech, return_tensors="pt", sampling_rate=16000).input_features |
| predicted_ids = whisper_model.generate(input_features, forced_decoder_ids=whisper_processor.get_decoder_prompt_ids(language=language, task="translate")) |
| transcription = whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0] |
| return transcription |
|
|
| def translate_and_summarize(text): |
| inputs = rag_tokenizer(text, return_tensors="pt") |
| input_ids = inputs["input_ids"] |
| attention_mask = inputs["attention_mask"] |
| outputs = rag_model.generate(input_ids=input_ids, attention_mask=attention_mask) |
| return rag_tokenizer.batch_decode(outputs, skip_special_tokens=True) |
|
|
| def extract_audio_from_video(video_path, output_audio_path): |
| video_clip = VideoFileClip(video_path) |
| audio_clip = video_clip.audio |
| if audio_clip is not None: |
| audio_clip.write_audiofile(output_audio_path) |
| return output_audio_path |
| else: |
| return None |
|
|
| st.title("Audio and Video Transcription & Summarization") |
|
|
| |
| st.header("Upload an Audio File") |
| audio_file = st.file_uploader("Choose an audio file...", type=["wav", "mp3", "m4a"]) |
|
|
| if audio_file is not None: |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file: |
| tmp_file.write(audio_file.getbuffer()) |
| audio_path = tmp_file.name |
| |
| st.audio(audio_file) |
| st.write("Transcribing audio...") |
| try: |
| transcription = transcribe_audio(audio_path) |
| st.write("Transcription:", transcription) |
| |
| st.write("Translating and summarizing...") |
| summary = translate_and_summarize(transcription) |
| st.write("Translated Summary:", summary) |
| except Exception as e: |
| st.error(f"An error occurred: {e}") |
|
|
| |
| st.header("Upload a Video File") |
| video_file = st.file_uploader("Choose a video file...", type=["mp4", "mkv", "avi", "mov"]) |
|
|
| if video_file is not None: |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_file: |
| tmp_file.write(video_file.getbuffer()) |
| video_path = tmp_file.name |
| |
| st.video(video_file) |
| st.write("Extracting audio from video...") |
| audio_path = extract_audio_from_video(video_path, tempfile.NamedTemporaryFile(delete=False, suffix=".wav").name) |
| |
| if audio_path is not None: |
| st.write("Transcribing audio...") |
| try: |
| transcription = transcribe_audio(audio_path) |
| st.write("Transcription:", transcription) |
| |
| st.write("Translating and summarizing...") |
| summary = translate_and_summarize(transcription) |
| st.write("Translated Summary:", summary) |
| except Exception as e: |
| st.error(f"An error occurred: {e}") |
| else: |
| st.write("No audio track found in the video file.") |
|
|
