Update app.py

app.py

@@ -4,42 +4,55 @@ import re
 
 client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
-def parse_srt(content):
-    """Parse the SRT file content, yielding text lines and skipping timestamps, line numbers, or blank lines."""
-    lines = content.split("\n")
-    for line in lines:
-        if not line.isdigit() and not re.match(r'\d{2}:\d{2}:\d{2},\d{3} --> \d{2}:\d{2}:\d{2},\d{3}', line) and line.strip():
-            yield line
+def translate_subtitles(srt_content, target_language):
+    """
+    Translate the subtitles in the SRT file content to the target language.
+    """
+    # Split the SRT content into blocks
+    blocks = srt_content.split('\n\n')
+    translated_blocks = []
+    for block in blocks:
+        if block.strip() == "":
+            continue
+        lines = block.split('\n')
+        if len(lines) >= 3:
+            index = lines[0]
+            time_range = lines[1]
+            subtitle_text = '\n'.join(lines[2:])
+            # Translate the subtitle text
+            translation = client.translation(subtitle_text, target_language=target_language)
+            translated_text = translation[0]['translation_text']
+            translated_blocks.append(f"{index}\n{time_range}\n{translated_text}")
+    return '\n\n'.join(translated_blocks)
 
-def translate_line(line, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
-    generate_kwargs = dict(
-        temperature=temperature,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        do_sample=True,
-        seed=42,
-    )
+def read_srt_file(file_path):
+    """
+    Read SRT file content.
+    """
+    with open(file_path, 'r', encoding='utf-8') as file:
+        return file.read()
 
-    stream = client.text_generation(f"Translate to English: {line}", **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-    for response in stream:
-        output += response.token.text
-    return output
+def save_translated_srt(content, output_path):
+    """
+    Save the translated subtitles to a new SRT file.
+    """
+    with open(output_path, 'w', encoding='utf-8') as file:
+        file.write(content)
 
-def translate_srt(file_info):
-    # Correctly access the file content from the file_info dictionary
-    content = file_info["content"].decode("utf-8")  # Decode content from bytes to string
-    translated_lines = []
-    for line in parse_srt(content):
-        translated_line = translate_line(line)
-        translated_lines.append(translated_line)
-    return "\n".join(translated_lines)
+def translate_srt_interface(srt_file, target_language):
+    """
+    Gradio interface function to translate SRT file content.
+    """
+    srt_content = srt_file.read()
+    translated_content = translate_subtitles(srt_content, target_language)
+    return translated_content
 
-gr.Interface(
-    fn=translate_srt,
-    inputs=gr.File(label="Upload SRT File"),
-    outputs=gr.Textbox(label="Translated Text"),
-    title="SRT Translator",
-    description="Upload an SRT file to translate its content line by line."
-).launch()
+iface = gr.Interface(
+    fn=translate_srt_interface,
+    inputs=[gr.File(file_type="srt"), gr.Textbox(label="Target Language Code")],
+    outputs="text",
+    title="SRT File Translator",
+    description="Translate SRT subtitle files to your desired language."
+)
+
+iface.launch()