Update app.py

app.py

@@ -1,168 +1,42 @@
-!pip install pysrt
-
-import google.generativeai as genai
-import pysrt
-import re
+# app.py
+import gradio as gr
 import os
-import csv
 from google.colab import files
+from core import process_files
 
-# 1. 配置 Gemini API - 使用者輸入金鑰
-while True:
-    api_key = input("請輸入您的 Google Gemini API 金鑰: ")
-    if api_key:
-      try:
-        genai.configure(api_key=api_key)
-        model = genai.GenerativeModel('gemini-1.5-flash') # 測試金鑰是否正確
-        print("API 設定完成！")
-        break
-      except Exception as e:
-        print(f"API 金鑰錯誤，請重新輸入: {e}")
-    else:
-        print("請輸入 API 金鑰")
-
-
-# 解析 SRT 檔案
-def parse_srt(srt_path):
-    try:
-        subs = pysrt.open(srt_path, encoding='utf-8')
-        srt_data = {}
-        for sub in subs:
-            srt_data[sub.index] = {"time": str(sub.start) + " --> " + str(sub.end), "text": sub.text}
-        return srt_data
-    except Exception as e:
-        print(f"解析字幕檔案錯誤: {e}")
-        return None
-
-# 2. 使用者介面上傳逐字稿檔案
-print("請上傳逐字稿檔案 (TXT 格式)")
-uploaded_transcript_file = files.upload()
-if uploaded_transcript_file:
-    test_transcript_file = list(uploaded_transcript_file.keys())[0]
-else:
-    print("沒有上傳任何檔案，請重新執行")
-    exit()
-
-# 檢查測試用逐字稿檔案是否存在
-if not os.path.exists(test_transcript_file):
-    print(f"錯誤: 找不到逐字稿檔案: {test_transcript_file}，請確認您已上傳該檔案")
-else:
-    with open(test_transcript_file, 'r', encoding='utf-8') as f:
+def correct_subtitle(api_key, transcript_file, oral_file, batch_size):
+    
+    with open(transcript_file.name, 'r', encoding='utf-8') as f:
         transcript_content = f.read()
 
-    # 預處理逐字稿
-    def preprocess_transcript(transcript):
-        # 移除 SSML 標籤
-        transcript = re.sub(r'<[^>]*?>', '', transcript)
-        # 移除多個連續空格
-        transcript = re.sub(r'\s+', ' ', transcript)
-        # 移除所有換行符號
-        transcript = transcript.replace('\n','')
-        # 移除文字前後的空白
-        transcript = transcript.strip()
-        return transcript
-    # 使用預處理函數處理逐字稿
-    transcript_content = preprocess_transcript(transcript_content)
-
-    # 3. 使用者介面上傳口語稿檔案
-    print("請上傳口語稿檔案 (SRT 格式)")
-    uploaded_oral_file = files.upload()
-    if uploaded_oral_file:
-        test_srt_file = list(uploaded_oral_file.keys())[0]
-    else:
-        print("沒有上傳任何檔案，請重新執行")
-        exit()
+    error, new_srt, all_reports = process_files(api_key, transcript_file.name, oral_file.name, batch_size)
 
-    # 檢查測試用 srt 檔案是否存在
-    if not os.path.exists(test_srt_file):
-        print(f"錯誤: 找不到口語稿檔案: {test_srt_file}，請確認您已上傳該檔案")
+    if error:
+      return error, None, None
     else:
-        srt_data = parse_srt(test_srt_file)
-
-        if srt_data is None:
-            print("錯誤: 解析口語稿檔案失敗，請確認檔案內容")
-        else:
-            # 4. 使用者介面輸入批次大小
-            while True:
-                try:
-                    batch_size = int(input("請輸入批次大小 (請輸入整數): "))
-                    if batch_size > 0:
-                        break
-                    else:
-                        print("錯誤：批次大小必須是正整數。")
-                except ValueError:
-                    print("錯誤：請輸入有效的整數。")
-
-            # 儲存所有報告
-            all_reports = []
-
-            # 處理每個批次
-            keys = list(srt_data.keys()) # 取得編號
-
-            for i in range(0, len(keys), batch_size):
-                batch_keys = keys[i:i + batch_size]
-                batch_data = {key: srt_data[key] for key in batch_keys} # 取得該批次的資料
-                subtitle_content = "\n".join([f"編號{index}:{data['text']}" for index, data in batch_data.items()])
-                prompt = (
-                     "請扮演一位專業的文字編輯，比對以下字幕檔和逐字稿，並以逐字稿的內容為準，修正字幕檔中因語音轉文字造成的錯別字、漏字、多字和詞語誤用等問題，請勿修改語氣和修辭，直接輸出修正後的字幕檔。\n"
-                    "接著，請將**當前批次**所有修改的內容整理成一個修改清單，並放在所有字幕內容之後。\n"
-                    "請使用以下格式輸出修改清單： \n"
-                    "編號: 原文: 修正後\n"
-                    "請使用 `<<<分隔符號>>>` 來分隔字幕與修改清單。\n"
-                     "請按照以下步驟進行：\n"
-                    "1. 比對字幕檔和逐字稿。\n"
-                    "2. 以逐字稿的內容為準，修正字幕檔中因語音轉文字造成的錯���字、漏字、多字和詞語誤用、以及標點符號漏缺等問題。\n"
-                    "3. 必須嚴格保持原始字幕的時間戳和斷句結構，不要拆分或合併任何段落，不要新增或移除任何字幕段落，直接輸出修正後的字幕檔。\n"
-                    "4. 讀取整份逐字稿時，請忽略SSML標籤，直接比對文字與標點符號，請勿修改任何語義。\n"
-                     "字幕檔內容：\n"
-                    f"{subtitle_content}\n\n"
-                    "逐字稿內容：\n"
-                     f"{transcript_content}\n\n"
-                   )
-                try:
-                    model = genai.GenerativeModel('gemini-1.5-flash')
-                    response = model.generate_content(prompt)
-                    corrected_subtitle = response.text
-                    print(f"第 {i // batch_size + 1} 批次 Gemini 模型的回應：")
-                    print(corrected_subtitle)
-
-                    # 使用 re.split 分割字幕和報告
-                    parts = re.split(r'<<<分隔符號>>>', corrected_subtitle, maxsplit=1) # 只分割一次
-                    corrected_lines = parts[0].split('\n')  # 取得字幕文字
-                    report = ""
-                    if len(parts) > 1:
-                        report = parts[1] # 取得修改清單文字
-                        report_lines = report.strip().split('\n')
-                        all_reports.extend(report_lines)
-
-                    for line in corrected_lines:
-                       match = re.match(r'編號(\d+):(.*)', line)
-                       if match:
-                           index = int(match.group(1))
-                           corrected_text = match.group(2).strip()
-                           srt_data[index]['text'] = corrected_text
-                except Exception as e:
-                    print(f"第 {i // batch_size + 1} 批次修正過程失敗：{e}")
-
-            # 寫回 SRT 檔案
-            new_subs = [] # 將 new_subs = [] 移到此
-            for index, data in srt_data.items():
-                time_str = data['time']
-                start_str, end_str = time_str.split(" --> ")
-                start = pysrt.srtitem.SubRipTime.from_string(start_str)
-                end = pysrt.srtitem.SubRipTime.from_string(end_str)
-                new_subs.append(pysrt.SubRipItem(index=index, start=start, end=end, text=data['text']))
-
-             # 將 new_subs 轉換為 SubRipFile 物件
-            new_srt = pysrt.SubRipFile(items=new_subs)
-
-            corrected_file_name = "corrected_subtitle.srt"
-            new_srt.save(corrected_file_name, encoding='utf-8')
-            files.download(corrected_file_name)
-
-            # 合併修改清單並下載
-            report_text = "\n".join(all_reports)
-            report_file_name = "report.txt"
-            with open(report_file_name, 'w', encoding='utf-8') as f:
-                f.write(report_text)
-            files.download(report_file_name)
+       corrected_file_name = "corrected_subtitle.srt"
+       new_srt.save(corrected_file_name, encoding='utf-8')
+       report_text = "\n".join(all_reports)
+       report_file_name = "report.txt"
+       with open(report_file_name, 'w', encoding='utf-8') as f:
+           f.write(report_text)
+       return "完成",corrected_file_name, report_file_name
+
+
+iface = gr.Interface(
+    fn=correct_subtitle,
+    inputs=[
+        gr.Textbox(lines=1, placeholder="請輸入您的 Google Gemini API 金鑰"),
+        gr.File(label="請上傳逐字稿檔案 (TXT 格式)"),
+        gr.File(label="請上傳口語稿檔案 (SRT 格式)"),
+        gr.Number(label="請輸入批次大小 (請輸入整數)")
+        ],
+    outputs=[
+        gr.Textbox(label="狀態"),
+        gr.File(label="修正後的字幕檔案"),
+        gr.File(label="修改清單")
+    ],
+    title="字幕修正工具"
+)
+
+iface.launch()