import gradio as gr
from faster_whisper import WhisperModel
from pydub import AudioSegment
import os
import tempfile
from transformers import pipeline

# הגדרת המודל לתמלול
model = WhisperModel("ivrit-ai/faster-whisper-v2-d4")

# הגדרת pipeline לסיכום עם מודל מותאם לעברית
summarizer = pipeline("summarization", model="yam-peleg/Hebrew-Mistral-7B-200K")

def transcribe_and_summarize(file_path):
    try:
        # בדיקה אם הקובץ הוא וידאו והמרת וידאו לאודיו במידת הצורך
        if file_path.endswith((".mp4", ".mov", ".avi", ".mkv")):
            audio_file = convert_video_to_audio(file_path)
        else:
            audio_file = file_path

        # תמלול האודיו
        segments, _ = model.transcribe(audio_file, language="he")
        transcript = " ".join([segment.text for segment in segments])

        # סיכום התמלול עם שימוש ב-max_new_tokens
        summary = summarizer(transcript, max_new_tokens=100)[0]["summary_text"]

        # מחיקת קובץ האודיו במידת הצורך (אם היה וידאו)
        if audio_file != file_path:
            os.remove(audio_file)

        return transcript, summary

    except Exception as e:
        return f"שגיאה בעיבוד הקובץ: {str(e)}", ""

def convert_video_to_audio(video_file):
    # יצירת קובץ אודיו זמני
    temp_audio = tempfile.mktemp(suffix=".wav")
    video = AudioSegment.from_file(video_file)
    video.export(temp_audio, format="wav")
    return temp_audio

# הגדרת ממשק Gradio
interface = gr.Interface(
    fn=transcribe_and_summarize,
    inputs=gr.Audio(type="filepath"),
    outputs=[
        gr.Textbox(label="תמלול"),
        gr.Textbox(label="סיכום")
        ],
    title="ממיר אודיו/וידאו לתמלול וסיכום",
    description="העלה קובץ אודיו או וידאו של מרצה וקבל תמלול מלא וסיכום קצר של התוכן."
)

if __name__ == "__main__":
    interface.launch()