import gradio as gr
import torch
import soundfile as sf
import librosa
from moviepy.editor import VideoFileClip
import os

# Load Whisper base model and processor
whisper_model_name = "openai/whisper-base"
whisper_processor = WhisperProcessor.from_pretrained(whisper_model_name)
whisper_model = WhisperForConditionalGeneration.from_pretrained(whisper_model_name)

# Load RAG sequence model and tokenizer
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)
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

def transcribe_audio_interface(audio_file):
    audio_path = os.path.join("/tmp", audio_file.name)
    with open(audio_path, "wb") as f:
        f.write(audio_file.getvalue())
    transcription = transcribe_audio(audio_path)
    return transcription

def summarize_text_interface(text):
    summary = translate_and_summarize(text)
    return summary

def summarize_video_interface(video_file):
    video_path = os.path.join("/tmp", video_file.name)
    with open(video_path, "wb") as f:
        f.write(video_file.getvalue())
    audio_path = extract_audio_from_video(video_path, "/tmp/extracted_audio.wav")
    if audio_path is not None:
        transcription = transcribe_audio(audio_path)
        summary = translate_and_summarize(transcription)
        return summary
    else:
        return "No audio track found in the video file."

# Create interfaces
audio_transcription_interface = gr.Interface(transcribe_audio_interface, inputs="audio", outputs="text", title="Audio Transcription")
text_summarization_interface = gr.Interface(summarize_text_interface, inputs="text", outputs="text", title="Text Summarization")
video_summarization_interface = gr.Interface(summarize_video_interface, inputs="video", outputs="text", title="Video Summarization")

# Launch the interfaces
audio_transcription_interface.launch()
text_summarization_interface.launch()
video_summarization_interface.launch()