Improve EPUB restoration, book terminology and Calibre QA workflow
This commit is contained in:
Kai Tan
@@ -0,0 +1,233 @@
|
||||
import asyncio
|
||||
import json
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
import pytest
|
||||
from bs4 import BeautifulSoup
|
||||
from ebooklib import epub
|
||||
from src.common.data_model import ManifestEntry, BookProfile, BookStructure
|
||||
from src.preprocessing.format_extractor import FormatExtractor
|
||||
from src.preprocessing.epub_cleaner import EpubCleaner
|
||||
from src.preprocessing.text_extractor import FineGrainedExtractor
|
||||
from src.assembly.format_restorer import FormatRestorer
|
||||
from src.assembly.restoration_engine import RestorationEngine
|
||||
from src.assembly.backfiller import BackfillEngine
|
||||
from src.assembly.builder import BilingualBuilder
|
||||
from src.translation.translator_engine import Translator
|
||||
from src.translation.glossary import BookGlossary, matches
|
||||
from src.translation.manifest_manager import ManifestManager
|
||||
from src.translation.llm_client import LLMClient
|
||||
from src.qa.epub_audit import audit_epub, compare_audits
|
||||
|
||||
|
||||
def entry(html='<p id="p1">Masayoshi Son <em>built</em> a company.</p>'):
|
||||
_, text, mapping, _, _ = FormatExtractor().extract(html)
|
||||
return ManifestEntry(entry_id='Text/ch.xhtml#p1', element_id='p1', file_path='Text/ch.xhtml',
|
||||
original_text=text, placeholders=mapping)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('html', [
|
||||
'<p>Hello <em><b>world</b></em>.</p>',
|
||||
'<p>The angle φ = 10, and φ1φ is literal mathematical text.</p>',
|
||||
'<p>A & B <span title="a > b">example</span>.</p>',
|
||||
'<p>A<a epub:type="noteref" href="notes.xhtml#n1">[1]</a> note.</p>',
|
||||
'<p>日本語と中文 <img src="x.png"/><br/>text.</p>',
|
||||
'<p>A <math xmlns="http://www.w3.org/1998/Math/MathML"><mi>x</mi></math> formula.</p>',
|
||||
'<p><span class="drop">T</span>he first word.</p>',
|
||||
])
|
||||
def test_dom_roundtrip(html):
|
||||
e = entry(html)
|
||||
restored, ok = FormatRestorer().restore(e.original_text, e.placeholders, original_text=e.original_text)
|
||||
assert ok
|
||||
before, after = BeautifulSoup(html, 'html.parser').p, BeautifulSoup('<p>'+restored+'</p>', 'html.parser').p
|
||||
assert before.get_text() == after.get_text()
|
||||
assert [(t.name, t.attrs) for t in before.find_all(True)] == [(t.name, t.attrs) for t in after.find_all(True)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize('text', [
|
||||
'φ1φφ2φ中文φ/1φφ/2φ', # crossed
|
||||
'φ1φφ2φ中文φ/2φφ/1φφ1φ', # duplicate
|
||||
'φ1φ中文φ/1φ', # missing pair
|
||||
'φ1φφ2φ中文φ/2φφ/1φφ99φ', # unknown
|
||||
'φ1φφ2φ中文φ/2φφ/1φφbrokenφ',
|
||||
])
|
||||
def test_invalid_placeholders_never_restore(text):
|
||||
e = entry('<p>Hello <em><b>world</b></em></p>')
|
||||
assert FormatRestorer().restore(text, e.placeholders, original_text=e.original_text) == ('', False)
|
||||
|
||||
|
||||
def test_model_text_is_escaped():
|
||||
result, ok = FormatRestorer().restore('<script>alert(1)</script> A & B', {})
|
||||
assert ok and '<script>' not in result and '&' in result
|
||||
|
||||
|
||||
def test_failed_restore_discards_cached_html():
|
||||
e = entry(); e.translated_text = '孙正义创办公司。'; e.translated_html = '<em>broken cached output'
|
||||
assert asyncio.run(RestorationEngine().restore_entries([e])) == 0
|
||||
assert e.translated_html is None
|
||||
|
||||
|
||||
def make_book(tmp_path):
|
||||
path = tmp_path/'source.epub'
|
||||
book = epub.EpubBook(); book.set_identifier('qa'); book.set_title('QA'); book.set_language('en')
|
||||
chapter = epub.EpubHtml(uid='ch', file_name='Text/ch.xhtml', title='Chapter')
|
||||
chapter.content = '''<html><head><style>.special{color:red}</style></head><body>
|
||||
<h1 id="title">Chapter</h1><p id="p1" class="special" style="text-align:center">Masayoshi Son <em>built</em> a company.</p>
|
||||
<ul><li id="list">Lead text<p id="nested">Nested paragraph.</p><ul><li>Child item</li></ul></li></ul>
|
||||
<table><tr><td id="cell">Cell text</td></tr></table><div class="poem">Poem line<br/>Second line</div>
|
||||
<p id="note"><a id="ref" href="#p1">Reference</a></p></body></html>'''
|
||||
css = epub.EpubItem(uid='css', file_name='Styles/main.css', media_type='text/css', content=b'.special{font-style:italic}')
|
||||
unrelated = epub.EpubItem(uid='othercss', file_name='Styles/other.css', media_type='text/css', content=b'p{color:blue}')
|
||||
chapter.add_link(href='../Styles/main.css', rel='stylesheet', type='text/css')
|
||||
for item in [chapter,css,unrelated,epub.EpubNav(),epub.EpubNcx()]: book.add_item(item)
|
||||
book.toc = [epub.Link('Text/ch.xhtml#p1','Chapter','toc1')]; book.spine = ['nav',chapter]
|
||||
epub.write_epub(str(path),book)
|
||||
return path
|
||||
|
||||
|
||||
def prepare(tmp_path):
|
||||
path = make_book(tmp_path); work = tmp_path/'work'
|
||||
structure = BookStructure.load(EpubCleaner(path,work).clean())
|
||||
return path,work,structure,FineGrainedExtractor().extract(structure)
|
||||
|
||||
|
||||
def test_targets_are_disjoint_and_ids_stable(tmp_path):
|
||||
path,work,structure,entries = prepare(tmp_path)
|
||||
assert sum('Nested paragraph' in e.original_text for e in entries) == 1
|
||||
assert any('Lead text' in e.original_text for e in entries)
|
||||
assert any('Cell text' in e.original_text for e in entries)
|
||||
assert any('Poem line' in e.original_text for e in entries)
|
||||
second = BookStructure.load(EpubCleaner(path,tmp_path/'work2').clean())
|
||||
assert [e.entry_id for e in entries] == [e.entry_id for e in FineGrainedExtractor().extract(second)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize('mode', ['bilingual','target_only'])
|
||||
def test_package_keeps_assets_css_and_anchors(tmp_path, mode):
|
||||
path,work,structure,entries = prepare(tmp_path)
|
||||
for e in entries: e.translated_text = e.original_text.replace('Masayoshi Son', '孙正义')
|
||||
assert asyncio.run(RestorationEngine().restore_entries(entries)) == len(entries)
|
||||
asyncio.run(BackfillEngine().backfill(structure,entries,mode))
|
||||
output = tmp_path/'translated.epub'; BilingualBuilder(work,path).build(structure,output)
|
||||
assert not compare_audits(audit_epub(path),audit_epub(output))
|
||||
with zipfile.ZipFile(path) as original, zipfile.ZipFile(output) as translated:
|
||||
for name in original.namelist():
|
||||
if name.endswith(('.css','.ncx')) or name.endswith('nav.xhtml'):
|
||||
assert original.read(name) == translated.read(name)
|
||||
name=next(n for n in translated.namelist() if n.endswith('Text/ch.xhtml'))
|
||||
soup=BeautifulSoup(translated.read(name),'html.parser')
|
||||
links=[x['href'] for x in soup.find_all('link')]
|
||||
assert '../Styles/main.css' in links and '../Styles/other.css' not in links
|
||||
assert soup.find(id='p1') is not None
|
||||
if mode=='bilingual':
|
||||
assert not any(x.parent.name == 'ul' for x in soup.select('.translation'))
|
||||
# Repeated backfill must not duplicate translations.
|
||||
before=structure.resources['ch'].content
|
||||
asyncio.run(BackfillEngine().backfill(structure,entries,mode))
|
||||
assert before==structure.resources['ch'].content
|
||||
|
||||
|
||||
def test_resume_reconciles_changed_source(tmp_path):
|
||||
manager=ManifestManager(tmp_path/'manifest.json'); e=entry(); manager.add_entries([e]); e.translated_text='旧译'
|
||||
changed=entry('<p>Different source</p>'); manager.add_entries([changed]); manager.save(); manager.load()
|
||||
assert len(manager.entries)==1 and not manager.entries[0].translated_text
|
||||
|
||||
|
||||
class FakeLLM:
|
||||
model='test'
|
||||
def __init__(self): self.calls=[]
|
||||
async def translate_chunk(self, entries, **kwargs):
|
||||
self.calls.append((entries[:], kwargs))
|
||||
return {e.entry_id: ('坏译文' if len(self.calls)==1 else e.original_text.replace('Masayoshi Son','孙正义')) for e in entries}
|
||||
|
||||
|
||||
def test_selective_retry_and_frozen_glossary():
|
||||
e=entry(); llm=FakeLLM(); glossary=BookGlossary([{'source':'Masayoshi Son','target':'孙正义'}])
|
||||
translator=Translator(llm,glossary=glossary)
|
||||
asyncio.run(translator.translate([e], BookProfile(title='Test',author='Test')))
|
||||
assert len(llm.calls)==2 and llm.calls[0][1]['glossary']=={'Masayoshi Son':'孙正义'}
|
||||
assert e.translated_text and not e.validation_errors
|
||||
asyncio.run(translator.translate([e], BookProfile(title='Test',author='Test')))
|
||||
assert len(llm.calls)==2
|
||||
|
||||
|
||||
def test_api_error_stays_pending():
|
||||
class Broken:
|
||||
model='bad'
|
||||
async def translate_chunk(self,*args,**kwargs): raise RuntimeError('network down')
|
||||
e=entry(); asyncio.run(Translator(Broken(),validation_retries=1).translate([e], None))
|
||||
assert not e.translated_text and e.validation_errors
|
||||
|
||||
|
||||
def test_duplicate_response_ids_are_not_accepted():
|
||||
client=object.__new__(LLMClient)
|
||||
assert client._parse_short_id_response('#1: 甲\n#1: 乙\n#2: 丙',{'#1':'one','#2':'two'})=={'two':'丙'}
|
||||
|
||||
|
||||
def test_glossary_alias_conflicts_and_word_boundaries():
|
||||
assert matches('Son','Son met Masa') and not matches('Son','Sony grandson')
|
||||
with pytest.raises(ValueError):
|
||||
BookGlossary([{'source':'Masayoshi Son','target':'孙正义','aliases':['Masa']}, {'source':'Masa','target':'马萨'}])
|
||||
glossary=BookGlossary([{'source':'Masayoshi Son','target':'孙正义','aliases':['Masa']}])
|
||||
assert glossary.validate('Masa built a company','马萨创办了公司')
|
||||
assert not glossary.validate('Masa built a company','孙正义创办了公司')
|
||||
|
||||
|
||||
def test_html5_repairs_broken_markup_without_duplicate_ids(tmp_path):
|
||||
cleaner=EpubCleaner(tmp_path/'unused.epub',tmp_path/'work')
|
||||
content=cleaner._clean_html('<html><body><p id="x">One<p id="x">Two <b>bold</p><div>Three</div></body></html>','ch.xhtml')
|
||||
soup=BeautifulSoup(content,'html.parser')
|
||||
ids=[t['id'] for t in soup.find_all(id=True)]
|
||||
assert len(ids)==len(set(ids))
|
||||
assert 'One' in soup.get_text() and 'Two' in soup.get_text() and 'Three' in soup.get_text()
|
||||
|
||||
|
||||
def test_surname_family_context_is_not_person():
|
||||
glossary=BookGlossary([{'source':'Masayoshi Son','target':'孙正义','aliases':['Son','Masa']}])
|
||||
e=entry('<p>The Son family settled in Japan.</p>')
|
||||
assert not glossary.relevant([e])
|
||||
assert not glossary.validate(e.original_text,'孙家在日本定居。')
|
||||
assert glossary.validate('Son met Masa.','马萨在会面。')
|
||||
specific=BookGlossary([{'source':'Masayoshi Son','target':'孙正义','aliases':['Son']},
|
||||
{'source':'Son family','target':'孙家'}])
|
||||
assert not specific.validate('The Son family settled.','孙家定居了。')
|
||||
|
||||
|
||||
def test_backfill_rejects_incomplete_entries_without_changing_language(tmp_path):
|
||||
_,_,structure,entries=prepare(tmp_path)
|
||||
entries[0].translated_text='中文'
|
||||
with pytest.raises(ValueError,match='unvalidated'):
|
||||
asyncio.run(BackfillEngine().backfill(structure,entries,'target_only'))
|
||||
assert structure.metadata.language=='en'
|
||||
asyncio.run(BackfillEngine().backfill(structure,entries,'target_only',allow_partial=True))
|
||||
assert structure.metadata.language=='en'
|
||||
|
||||
|
||||
def test_explicit_glossary_context_exclusion():
|
||||
glossary=BookGlossary([{'source':'Apple','target':'苹果公司','exclude_contexts':['apple pie']}])
|
||||
assert not glossary.validate('An apple pie.','一个苹果派。')
|
||||
assert glossary.validate('Apple launched a product.','苹果发布了产品。')
|
||||
|
||||
|
||||
def test_mismatched_nested_response_id_is_not_assigned():
|
||||
client=object.__new__(LLMClient)
|
||||
assert client._parse_short_id_response('#1: #2: 错译',{'#1':'one','#2':'two'})=={}
|
||||
|
||||
|
||||
@pytest.mark.parametrize('finish_reason', ['stop', 'length'])
|
||||
def test_output_budget_and_truncation_handling(finish_reason):
|
||||
from types import SimpleNamespace
|
||||
from src.translation.llm_client import RateLimiter
|
||||
client=object.__new__(LLMClient)
|
||||
client.model='test';client.max_output_tokens=8192;client.temperature=None
|
||||
client.rate_limiter=RateLimiter(0,1)
|
||||
async def create(**kwargs):
|
||||
assert kwargs['max_completion_tokens']==8192
|
||||
assert 'temperature' not in kwargs
|
||||
return SimpleNamespace(choices=[SimpleNamespace(finish_reason=finish_reason,message=SimpleNamespace(content='译文'))])
|
||||
client.client=SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create)))
|
||||
call=client._make_request.__wrapped__(client,'system','source')
|
||||
if finish_reason=='stop':
|
||||
assert asyncio.run(call)=='译文'
|
||||
else:
|
||||
with pytest.raises(ValueError,match='Incomplete'):
|
||||
asyncio.run(call)
|
||||
Reference in New Issue
Block a user