Update app.py

app.py

@@ -3,6 +3,7 @@ import pandas as pd
 import gradio as gr
 from transformers import MarianMTModel, MarianTokenizer
 import io
+import pysrt
 
 # Fetch and parse language options
 url = "https://huggingface.co/Lenylvt/LanguageISO/resolve/main/iso.md"
@@ -35,21 +36,23 @@ def translate_text(text, source_language_code, target_language_code):
     
     return translated_text
 
-def translate_srt(file_info, source_language_code, target_language_code):
-    srt_content = file_info["content"].decode("utf-8")
-    translated_lines = []
-    for line in srt_content.split("\n"):
-        if line.isdigit() or "-->" in line:
-            translated_lines.append(line)
-        elif line.strip() == "":
-            translated_lines.append(line)
-        else:
-            translated_text = translate_text(line, source_language_code, target_language_code)
-            translated_lines.append(translated_text)
-    translated_srt_content = "\n".join(translated_lines)
+def translate_srt(file_content, source_language_code, target_language_code):
+    # Load the SRT content using pysrt and io.BytesIO
+    subs = pysrt.from_string(file_content.decode("utf-8"))
+
+    # Translate each subtitle
+    for sub in subs:
+        translated_text = translate_text(sub.text, source_language_code, target_language_code)
+        sub.text = translated_text
+    
+    # Convert the subtitles back to an SRT string
+    srt_content = '\n'.join(subs.to_string())
+
+    # Save the translated SRT content to a temporary file
     output_path = "/mnt/data/translated_srt.srt"
     with open(output_path, "w", encoding="utf-8") as file:
-        file.write(translated_srt_content)
+        file.write(srt_content)
+    
     return output_path
 
 source_language_dropdown = gr.Dropdown(choices=language_options, label="Source Language")