Spaces:
Running
Running
from transformers import MarianMTModel, MarianTokenizer | |
from tqdm import tqdm | |
import os | |
import re | |
import argparse | |
# Load Model and Tokenizer | |
model_name = "Helsinki-NLP/opus-mt-en-es" | |
tokenizer = MarianTokenizer.from_pretrained(model_name) | |
model = MarianMTModel.from_pretrained(model_name) | |
# Extract & separate timestamp and text | |
def extract_timestamp_and_text(line): | |
match = re.match(r'\[(\d+\.\d+\-\d+\.\d+)\]\s+(.*)', line) | |
if match: | |
return match.group(1), match.group(2) | |
return '', line | |
# Translate text | |
def translate_text(text): | |
lines = text.split('\n') | |
translated_lines = [] | |
for line in tqdm(lines, desc="Translating lines", leave=False): | |
if not line.strip(): | |
translated_lines.append('') | |
continue | |
timestamp, line_text = extract_timestamp_and_text(line) | |
if line_text.strip(): | |
model_inputs = tokenizer(line_text, return_tensors="pt", truncation=True, padding="longest") | |
translated = model.generate(**model_inputs) | |
translated_text = tokenizer.batch_decode(translated, skip_special_tokens=True)[0] | |
translated_line = f'[{timestamp}] {translated_text}' | |
else: | |
translated_line = f'[{timestamp}]' | |
translated_lines.append(translated_line) | |
return '\n'.join(translated_lines) | |
# Main function to translate a file | |
def translate_file(src_file_path, dst_file_path): | |
try: | |
with open(src_file_path, 'r') as file: | |
english_text = file.read() | |
spanish_text = translate_text(english_text) | |
with open(dst_file_path, 'w') as file: | |
file.write(spanish_text) | |
print(f"Translation completed: {dst_file_path}") | |
except Exception as e: | |
print(f"Error processing file: {e}") | |
if __name__ == "__main__": | |
parser = argparse.ArgumentParser(description="Translate English text to Spanish") | |
parser.add_argument("src_file_path", help="Path to the source file with English text") | |
parser.add_argument("dst_file_path", help="Path to save the translated Spanish text") | |
args = parser.parse_args() | |
translate_file(args.src_file_path, args.dst_file_path) | |