Update app.py

app.py

@@ -1,63 +1,65 @@
+import io
+import requests
 import streamlit as st
 import pandas as pd
 import pysrt
 from transformers import MarianMTModel, MarianTokenizer
 
-# Fetch and parse language options from the provided URL
-url = "https://huggingface.co/Lenylvt/LanguageISO/resolve/main/iso.md"
-df = pd.read_csv(url, 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()
-
-# Prepare language options for the dropdown
-language_options = df.set_index('ISO 639-1')[['Language Name']].to_dict()['Language Name']
+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:
+        st.error("Failed to fetch language options. Please try again later.")
+        return []
 
 def translate_text(text, source_language_code, target_language_code):
-    # Construct model name using ISO 639-1 codes
     model_name = f"Helsinki-NLP/opus-mt-{source_language_code}-{target_language_code}"
-
-    # Check if source and target languages are the same, which is not supported for translation
     if source_language_code == target_language_code:
         return "Translation between the same languages is not supported."
-
-    # Load tokenizer and model
     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)}"
-
-    # Translate text
-    translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512))
-    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
     
-    return translated_text
+    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):
-    # Load SRT file
     subs = pysrt.open(input_file.name)
-
-    # Initialize an empty list to store translated subtitles
     translated_subs = []
-
-    # Translate each subtitle
+    progress_bar = st.progress(0)
     for idx, sub in enumerate(subs):
         translated_text = translate_text(sub.text, source_language_code, target_language_code)
-        # Construct the translated subtitle with timestamp and line number
         translated_sub = pysrt.SubRipItem(index=idx+1, start=sub.start, end=sub.end, text=translated_text)
         translated_subs.append(translated_sub)
-
-    # Save translated subtitles to a new SRT file
+        progress_bar.progress((idx + 1) / len(subs))
     translated_file = pysrt.SubRipFile(translated_subs)
     translated_srt_path = input_file.name.replace(".srt", f"_{target_language_code}.srt")
     translated_file.save(translated_srt_path)
+    progress_bar.empty()
     return translated_srt_path
 
 st.title("SRT Translator")
 st.write("Translate subtitles from one language to another.")
 
-source_language_code = st.selectbox("Select Source Language", options=list(language_options.keys()))
-target_language_code = st.selectbox("Select Target Language", options=list(language_options.keys()))
+# Fetch language options
+url = "https://huggingface.co/Lenylvt/LanguageISO/resolve/main/iso.md"
+language_options = fetch_languages(url)
+
+source_language_code = st.selectbox("Select Source Language", options=language_options)
+target_language_code = st.selectbox("Select Target Language", options=language_options)
 
 file_input = st.file_uploader("Upload SRT File", type=["srt"])