SubtitleFixer / core.bak
hlyangster's picture
Update core.bak
afbf047 verified
# core.py
import google.generativeai as genai
import pysrt
import re
import os
# 解析 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
# 預處理逐字稿
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
def process_files(api_key, test_transcript_file, test_srt_file, batch_size):
# 1. 配置 Gemini API
try:
genai.configure(api_key=api_key)
model = genai.GenerativeModel('gemini-2.0-flash-exp') # 測試金鑰是否正確
except Exception as e:
return f"API 金鑰錯誤,請檢查設定: {e}", None, None
# 檢查測試用逐字稿檔案是否存在
if not os.path.exists(test_transcript_file):
return f"錯誤: 找不到逐字稿檔案: {test_transcript_file},請確認您已上傳該檔案", None, None
else:
with open(test_transcript_file, 'r', encoding='utf-8') as f:
transcript_content = f.read()
transcript_content = preprocess_transcript(transcript_content)
# 檢查測試用 srt 檔案是否存在
if not os.path.exists(test_srt_file):
return f"錯誤: 找不到口語稿檔案: {test_srt_file},請確認您已上傳該檔案", None, None
else:
srt_data = parse_srt(test_srt_file)
if srt_data is None:
return "錯誤: 解析口語稿檔案失敗,請確認檔案內容", None, None
# 儲存所有報告
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:
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:
return f"第 {i // batch_size + 1} 批次修正過程失敗:{e}", None, None
# 寫回 SRT 檔案
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)
return None, new_srt, all_reports