app.py

import gradio as gr
from huggingface_hub import InferenceClient
import re

# Initialize the InferenceClient with the Mixtral model
client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")

def translate_srt(file, target_language):
    # Read the content of the SRT file
    srt_content = file.read().decode("utf-8")
    lines = srt_content.split('\n')
    
    translated_lines = []
    for line in lines:
        # Check if the line is a timestamp or a subtitle number
        if re.match(r"^\d+$", line) or re.match(r"^\d{2}:\d{2}:\d{2},\d{3} --> \d{2}:\d{2}:\d{2},\d{3}$", line):
            translated_lines.append(line)  # Copy timestamps and numbers directly
        elif line.strip() == "":
            translated_lines.append(line)  # Preserve empty lines for formatting
        else:
            # Translate the text line
            response = client(inputs=line, parameters={"target_language": target_language})
            translated_lines.append(response[0]["generated_text"])
    
    # Join the translated lines back into a single string
    translated_srt_content = "\n".join(translated_lines)
    return translated_srt_content

# Gradio interface
iface = gr.Interface(
    fn=translate_srt,
    inputs=[gr.inputs.File(label="Upload SRT File"), gr.inputs.Dropdown(["fr", "en", "es", "de", "it", "pt"], label="Target Language")],
    outputs="text",
    title="SRT File Translator",
    description="Translate SRT files to the selected language using Mixtral model."
)

# Launch the Gradio app
iface.launch()