Spaces:
Running
Running
| import io | |
| import requests | |
| import streamlit as st | |
| import pandas as pd | |
| import pysrt | |
| from transformers import MarianMTModel, MarianTokenizer | |
| import tempfile | |
| def fetch_languages(url): | |
| response = requests.get(url) | |
| if response.status_code == 200: | |
| # Convert bytes to a string using decode, then create a file-like object with io.StringIO | |
| csv_content = response.content.decode('utf-8') | |
| df = pd.read_csv(io.StringIO(csv_content), delimiter="|", skiprows=2, header=None).dropna(axis=1, how='all') | |
| df.columns = ['ISO 639-1', 'ISO 639-2', 'Language Name', 'Native Name'] | |
| df['ISO 639-1'] = df['ISO 639-1'].str.strip() | |
| language_options = [(row['ISO 639-1'], f"{row['ISO 639-1']} - {row['Language Name']}") for index, row in df.iterrows()] | |
| return language_options | |
| else: | |
| return [] | |
| def translate_text(text, source_language_code, target_language_code): | |
| model_name = f"Helsinki-NLP/opus-mt-{source_language_code}-{target_language_code}" | |
| if source_language_code == target_language_code: | |
| return "Translation between the same languages is not supported." | |
| try: | |
| tokenizer = MarianTokenizer.from_pretrained(model_name) | |
| model = MarianMTModel.from_pretrained(model_name) | |
| except Exception as e: | |
| return f"Failed to load model for {source_language_code} to {target_language_code}: {str(e)}" | |
| translated_texts = [] | |
| for sentence in text.split("\n"): | |
| translated = model.generate(**tokenizer(sentence, return_tensors="pt", padding=True, truncation=True, max_length=512)) | |
| translated_text = tokenizer.decode(translated[0], skip_special_tokens=True) | |
| translated_texts.append(translated_text) | |
| return "\n".join(translated_texts) | |
| def translate_srt(input_file, source_language_code, target_language_code): | |
| subs = pysrt.open(input_file) | |
| translated_subs = [] | |
| progress_bar = st.progress(0) | |
| for idx, sub in enumerate(subs): | |
| translated_text = translate_text(sub.text, source_language_code, target_language_code) | |
| translated_sub = pysrt.SubRipItem(index=idx+1, start=sub.start, end=sub.end, text=translated_text) | |
| translated_subs.append(translated_sub) | |
| progress_bar.progress((idx + 1) / len(subs)) | |
| translated_file = pysrt.SubRipFile(translated_subs) | |
| # Use tempfile to create a temporary file path | |
| with tempfile.NamedTemporaryFile(suffix=".srt", delete=False) as tmp_file: | |
| translated_file.save(tmp_file.name) | |
| translated_srt_path = tmp_file.name | |
| progress_bar.empty() | |
| return translated_srt_path | |
| st.title("SRT Translator") | |
| st.write("Translate subtitles from one language to another.") | |
| # Fetch language options | |
| url = "https://huggingface.co/Lenylvt/LanguageISO/resolve/main/iso.md" | |
| language_options = fetch_languages(url) | |
| source_language = st.selectbox("Select Source Language", options=language_options, format_func=lambda x: x[1]) | |
| target_language = st.selectbox("Select Target Language", options=language_options, format_func=lambda x: x[1]) | |
| file_input = st.file_uploader("Upload SRT File", type=["srt"]) | |
| if file_input is not None: | |
| with tempfile.NamedTemporaryFile(suffix=".srt", delete=False) as temp_file: | |
| temp_file.write(file_input.read()) | |
| temp_file.seek(0) | |
| translated_srt_path = translate_srt(temp_file.name, source_language_code, target_language_code) | |
| st.success(f"Translation complete! You can download the translated SRT file from [here]({translated_srt_path})") | |