Spaces:
Sleeping
Sleeping
File size: 2,247 Bytes
f0ceee4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
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)
|