Update Milestone5API.py

Milestone5API.py

@@ -1,9 +1,10 @@
+import os
+import re
 from pytube import YouTube
 from youtube_transcript_api import YouTubeTranscriptApi
-import re
-import os
+from transformers import MarianMTModel, MarianTokenizer
 
-def download_video_transcript(video_url):
+def download_video_transcript(video_url, source_lang='en', target_lang='fr'):
     try:
         video_id = re.search(r"(?<=v=)[\w-]+", video_url)
         if video_id:
@@ -25,62 +26,31 @@ def download_video_transcript(video_url):
             print('Download completed!')
 
             transcript_list = YouTubeTranscriptApi.list_transcripts(video_id)
-            transcript = transcript_list.find_generated_transcript(['en']).fetch()
+            transcript = transcript_list.find_generated_transcript([source_lang]).fetch()
 
-            captions = ""
+            original_captions = ""
             for i, line in enumerate(transcript):
                 start_time = line['start']
                 formatted_time = f"{int(start_time // 60):02d}:{int(start_time % 60):02d}"
-                captions += f"{formatted_time} {line['text']}\n"
+                original_captions += f"{formatted_time} {line['text']}\n"
 
-            with open(f'{captions_path}/{modified_title}.txt', 'w', encoding='utf-8') as file:
-                file.write(captions)
+            original_filename = f'{captions_path}/{modified_title}_original.txt'
+            with open(original_filename, 'w', encoding='utf-8') as file:
+                file.write(original_captions)
             
+            # Translation part
+            translated_captions = translate_text_file(original_filename,
+                                                      f'{captions_path}/{modified_title}_translated.txt',
+                                                      source_lang=source_lang,
+                                                      target_lang=target_lang)
+
             # Delete video file after captions are saved
             os.remove(video_file)
             
-            return captions
+            return original_captions, translated_captions, original_filename, f'{captions_path}/{modified_title}_translated.txt'
         else:
             print("Video ID not found in URL.")
-            return None
+            return None, None, None, None
     except Exception as e:
         print("Error:", e)
-        return None
-from transformers import MarianMTModel, MarianTokenizer
-import os
-
-def translate_text_file(input_text_path, output_text_path, source_lang='en', target_lang='fr', batch_size=8):
-    model_name = "Helsinki-NLP/opus-mt-en-fr"
-
-    model = MarianMTModel.from_pretrained(model_name)
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-
-    def translate_batch(model, tokenizer, sentences):
-        sentences = [f"{source_lang}: {sentence}" for sentence in sentences]
-
-        input_ids = tokenizer(sentences, return_tensors="pt", padding=True, truncation=True)["input_ids"]
-        translation_ids = model.generate(input_ids)
-        translated_texts = tokenizer.batch_decode(translation_ids, skip_special_tokens=True)
-
-        return translated_texts
-
-    def read_text_from_file(file_path):
-        with open(file_path, 'r', encoding='utf-8') as file:
-            text = file.readlines()
-        return text
-
-    def write_text_to_file(file_path, translated_texts):
-        with open(file_path, 'w', encoding='utf-8') as file:
-            file.writelines([f"{line}\n" for line in translated_texts])
-
-    input_lines = read_text_from_file(input_text_path)
-    translated_lines = []
-
-    for i in range(0, len(input_lines), batch_size):
-        batch = input_lines[i:i + batch_size]
-        translated_batch = translate_batch(model, tokenizer, batch)
-        translated_lines.extend(translated_batch)
-
-    write_text_to_file(output_text_path, translated_lines)
-    return translated_lines
-
+        return None, None, None, None