Azhi75/kurdish-caption-generator

1import whisper2import gradio as gr3from transformers import MarianMTModel, MarianTokenizer, AutoTokenizer, AutoModelForSeq2SeqLM4import os5 6# Load models once7whisper_model = whisper.load_model("base")8sorani_model_name = "Helsinki-NLP/opus-mt-en-ku"9sorani_tokenizer = MarianTokenizer.from_pretrained(sorani_model_name)10sorani_model = MarianMTModel.from_pretrained(sorani_model_name)11 12kurmanji_model_name = "facebook/nllb-200-distilled-600M"13kurmanji_tokenizer = AutoTokenizer.from_pretrained(kurmanji_model_name)14kurmanji_model = AutoModelForSeq2SeqLM.from_pretrained(kurmanji_model_name)15 16# Format time to SRT style17def format_time(seconds):18    h = int(seconds // 3600)19    m = int((seconds % 3600) // 60)20    s = int(seconds % 60)21    ms = int((seconds - int(seconds)) * 1000)22    return f"{h:02}:{m:02}:{s:02},{ms:03}"23 24# Main processing function25def generate_captions(video):26    result = whisper_model.transcribe(video)27    segments = result['segments']28    english_lines = [seg['text'] for seg in segments]29 30    # Sorani Translation31    inputs = sorani_tokenizer(english_lines, return_tensors="pt", padding=True, truncation=True)32    outputs = sorani_model.generate(**inputs)33    sorani_lines = [sorani_tokenizer.decode(t, skip_special_tokens=True) for t in outputs]34 35    # Kurmanji Translation36    kurmanji_tokenizer.src_lang = "eng_Latn"37    inputs_km = kurmanji_tokenizer(english_lines, return_tensors="pt", padding=True, truncation=True)38    outputs_km = kurmanji_model.generate(**inputs_km, forced_bos_token_id=kurmanji_tokenizer.lang_code_to_id["kmr_Latn"])39    kurmanji_lines = [kurmanji_tokenizer.decode(t, skip_special_tokens=True) for t in outputs_km]40 41    # Create SRTs42    sorani_srt = ""43    kurmanji_srt = ""44    for i, seg in enumerate(segments):45        start = format_time(seg['start'])46        end = format_time(seg['end'])47        sorani_srt += f"{i+1}\n{start} --> {end}\n{sorani_lines[i]}\n\n"48        kurmanji_srt += f"{i+1}\n{start} --> {end}\n{kurmanji_lines[i]}\n\n"49 50    return sorani_srt, kurmanji_srt51 52# Gradio UI53interface = gr.Interface(54    fn=generate_captions,55    inputs=gr.Video(label="Upload your English video"),56    outputs=[57        gr.Textbox(label="Sorani Subtitles (.srt format)", lines=10),58        gr.Textbox(label="Kurmanji Subtitles (.srt format)", lines=10)59    ],60    title="Kurdish Subtitle Generator",61    description="Upload an English video. Get Sorani and Kurmanji Kurdish subtitles."62)63 64interface.launch()65