import streamlit as st
import pysrt
from transformers import MarianMTModel, MarianTokenizer

# تابع برای ترجمه متون
def translate_text(text, model, tokenizer):
    translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True))
    translated_text = [tokenizer.decode(t, skip_special_tokens=True) for t in translated]
    return translated_text[0]

# بارگذاری مدل و توکنایزر
model_name = 'Helsinki-NLP/opus-mt-en-fa'
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)

st.title("SRT Translator: English to Persian")
st.write("Upload an SRT file to translate it from English to Persian.")

uploaded_file = st.file_uploader("Choose an SRT file", type="srt")

if uploaded_file is not None:
    subs = pysrt.from_string(uploaded_file.read().decode('utf-8'))
    for sub in subs:
        sub.text = translate_text(sub.text, model, tokenizer)
    
    output_srt = subs.to_string()
    st.download_button(
        label="Download Translated SRT",
        data=output_srt,
        file_name="translated.srt",
        mime="text/srt"
    )
    st.text_area("Translated SRT Preview", output_srt, height=300)