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)