Lenylvt commited on
Commit
c64c3dc
β€’
1 Parent(s): 547eed9

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +66 -0
app.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import pandas as pd
3
+ import pysrt
4
+ from transformers import MarianMTModel, MarianTokenizer
5
+
6
+ # Fetch and parse language options from the provided URL
7
+ url = "https://huggingface.co/Lenylvt/LanguageISO/resolve/main/iso.md"
8
+ df = pd.read_csv(url, delimiter="|", skiprows=2, header=None).dropna(axis=1, how='all')
9
+ df.columns = ['ISO 639-1', 'ISO 639-2', 'Language Name', 'Native Name']
10
+ df['ISO 639-1'] = df['ISO 639-1'].str.strip()
11
+
12
+ # Prepare language options for the dropdown
13
+ language_options = df.set_index('ISO 639-1')[['Language Name']].to_dict()['Language Name']
14
+
15
+ def translate_text(text, source_language_code, target_language_code):
16
+ # Construct model name using ISO 639-1 codes
17
+ model_name = f"Helsinki-NLP/opus-mt-{source_language_code}-{target_language_code}"
18
+
19
+ # Check if source and target languages are the same, which is not supported for translation
20
+ if source_language_code == target_language_code:
21
+ return "Translation between the same languages is not supported."
22
+
23
+ # Load tokenizer and model
24
+ try:
25
+ tokenizer = MarianTokenizer.from_pretrained(model_name)
26
+ model = MarianMTModel.from_pretrained(model_name)
27
+ except Exception as e:
28
+ return f"Failed to load model for {source_language_code} to {target_language_code}: {str(e)}"
29
+
30
+ # Translate text
31
+ translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512))
32
+ translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
33
+
34
+ return translated_text
35
+
36
+ def translate_srt(input_file, source_language_code, target_language_code):
37
+ # Load SRT file
38
+ subs = pysrt.open(input_file.name)
39
+
40
+ # Initialize an empty list to store translated subtitles
41
+ translated_subs = []
42
+
43
+ # Translate each subtitle
44
+ for idx, sub in enumerate(subs):
45
+ translated_text = translate_text(sub.text, source_language_code, target_language_code)
46
+ # Construct the translated subtitle with timestamp and line number
47
+ translated_sub = pysrt.SubRipItem(index=idx+1, start=sub.start, end=sub.end, text=translated_text)
48
+ translated_subs.append(translated_sub)
49
+
50
+ # Save translated subtitles to a new SRT file
51
+ translated_file = pysrt.SubRipFile(translated_subs)
52
+ translated_srt_path = input_file.name.replace(".srt", f"_{target_language_code}.srt")
53
+ translated_file.save(translated_srt_path)
54
+ return translated_srt_path
55
+
56
+ st.title("SRT Translator")
57
+ st.write("Translate subtitles from one language to another.")
58
+
59
+ source_language_code = st.selectbox("Select Source Language", options=list(language_options.keys()))
60
+ target_language_code = st.selectbox("Select Target Language", options=list(language_options.keys()))
61
+
62
+ file_input = st.file_uploader("Upload SRT File", type=["srt"])
63
+
64
+ if file_input is not None:
65
+ translated_srt_path = translate_srt(file_input, source_language_code, target_language_code)
66
+ st.success(f"Translation complete! You can download the translated SRT file from [here]({translated_srt_path})")