feat: support single_translate for txt & srt (#265)

* feat: support single_translate for txt & srt

* style: black format

* chore: change func name
This commit is contained in:
zstone12 2023-05-07 21:22:46 +08:00 committed by GitHub
parent 60df5f2723
commit b5060b8746
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 24 additions and 8 deletions

View File

@ -43,6 +43,7 @@ class SRTBookLoader(BaseBookLoader):
self.test_num = test_num
self.accumulated_num = 1
self.blocks = []
self.single_translate = single_translate
self.resume = resume
self.bin_path = f"{Path(srt_name).parent}/.{Path(srt_name).stem}.temp.bin"
@ -76,6 +77,9 @@ class SRTBookLoader(BaseBookLoader):
def _get_block_text(self, block):
return f"{block['number']}\n{block['time']}\n{block['text']}"
def _get_block_except_text(self, block):
return f"{block['number']}\n{block['time']}"
def _concat_blocks(self, sliced_text: str, text: str):
return f"{sliced_text}\n\n{text}" if sliced_text else text
@ -210,12 +214,22 @@ class SRTBookLoader(BaseBookLoader):
for i, block in enumerate(translated_blocks):
text = block.get("text", "")
self.p_to_save.append(text)
if self.single_translate:
self.bilingual_result.append(
f"{self._get_block_except_text(self.blocks[begin + i])}\n{text}"
)
else:
self.bilingual_result.append(
f"{self._get_block_text(self.blocks[begin + i])}\n{text}"
)
else:
for i, block in enumerate(self.blocks[begin:end]):
text = self.p_to_save[begin + i]
if self.single_translate:
self.bilingual_result.append(
f"{self._get_block_except_text(self.blocks[begin + i])}\n{text}"
)
else:
self.bilingual_result.append(
f"{self._get_block_text(self.blocks[begin + i])}\n{text}"
)

View File

@ -33,6 +33,7 @@ class TXTBookLoader(BaseBookLoader):
self.bilingual_temp_result = []
self.test_num = test_num
self.batch_size = 10
self.single_translate = single_translate
try:
with open(f"{txt_name}", encoding="utf-8") as f:
@ -73,6 +74,7 @@ class TXTBookLoader(BaseBookLoader):
print(e)
raise Exception("Something is wrong when translate") from e
self.p_to_save.append(temp)
if not self.single_translate:
self.bilingual_result.append(batch_text)
self.bilingual_result.append(temp)
index += self.batch_size