Update app.py

app.py

@@ -1,62 +1,70 @@
-from huggingface_hub import InferenceClient
+import requests
+import pandas as pd
 import gradio as gr
+from transformers import MarianMTModel, MarianTokenizer
+import io
 
-client = InferenceClient(
-    "mistralai/Mixtral-8x7B-Instruct-v0.1"
-)
+# Fetch and parse language options
+url = "https://huggingface.co/Lenylvt/LanguageISO/resolve/main/iso.md"
+response = requests.get(url)
+df = pd.read_csv(io.StringIO(response.text), 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()
 
-def format_prompt(message, history, target_language):
-    prompt = "<s>"
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    prompt += f"[INST] {message} [/INST] [Translate to: {target_language}]"
-    return prompt
-
-def generate_from_srt(file_content, target_language, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
-    # Process the SRT file content as needed before using it as a prompt
-    # For example, extracting text and removing timestamps if necessary
-    # Directly using the file content for simplicity here
-    
-    temperature = float(temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
-    top_p = float(top_p)
-
-    generate_kwargs = dict(
-        temperature=temperature,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        do_sample=True,
-        seed=42,
-    )
-
-    formatted_prompt = format_prompt(file_content, [], target_language)
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-
-    for response in stream:
-        output += response.token.text
-    return output
-
-def handle_file(file_info, target_language):
-    # Directly use the file content if it's a text file
-    if isinstance(file_info, str):
-        file_content = file_info
-    else:
-        # If file_info is not a string, it might be a binary file
-        file_content = file_info.decode('utf-8')
+# Prepare language options for the dropdown
+language_options = [(row['ISO 639-1'], f"{row['ISO 639-1']} - {row['Language Name'].strip()}") for index, row in df.iterrows()]
+
+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
+    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 generate_from_srt(file_content, target_language)
+    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)
+    output_path = "/mnt/data/translated_srt.srt"
+    with open(output_path, "w", encoding="utf-8") as file:
+        file.write(translated_srt_content)
+    return output_path
+
+source_language_dropdown = gr.Dropdown(choices=language_options, label="Source Language")
+target_language_dropdown = gr.Dropdown(choices=language_options, label="Target Language")
 
 iface = gr.Interface(
-    fn=handle_file,
-    inputs=[gr.File(label="Upload SRT File"), gr.Textbox(label="Target Language", placeholder="Enter target language (e.g., Spanish, French)")],
-    outputs="text",
-    title="SRT File Translation",
-    description="Translate the content of SRT files to the specified language.",
-    concurrency_limit=20,
+    fn=translate_srt,
+    inputs=[
+        gr.File(label="Upload SRT File", type="file"),
+        source_language_dropdown,
+        target_language_dropdown
+    ],
+    outputs=gr.File(label="Download Translated SRT File"),
+    title="SRT Translator",
+    description="Translate SubRip Text (SRT) subtitle files. This tool uses models from the Language Technology Research Group at the University of Helsinki."
 )
 
-iface.launch()
+iface.launch()