TranslateBooksWithLLMs/test_epub_refinement_progress.py at main · hydropix/TranslateBooksWithLLMs · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
"""
Test script to verify EPUB progress tracking with refinement enabled.
"""

from src.core.epub.translation_metrics import TranslationMetrics

def test_epub_refinement_progress():
    """Test that EPUB refinement doubles the total work and progress is calculated correctly."""

    print("=" * 60)
    print("Testing EPUB TranslationMetrics with Refinement")
    print("=" * 60)

    # Create metrics with refinement enabled
    stats = TranslationMetrics()
    stats.total_chunks = 5  # 5 chunks to translate
    stats.enable_refinement = True
    stats.refinement_phase = False

    print(f"\nInitial state:")
    stats_dict = stats.to_dict()
    print(f"  Total chunks: {stats_dict['total_chunks']}")
    print(f"  Completed chunks: {stats_dict['completed_chunks']}")
    print(f"  Progress: {(stats_dict['completed_chunks'] / stats_dict['total_chunks']) * 100:.1f}%")

    # Complete phase 1 (translation)
    print("\n--- PHASE 1: TRANSLATION ---")
    for i in range(5):
        stats.successful_first_try = i + 1
        stats.record_processed()  # Mark chunk as processed
        stats_dict = stats.to_dict()
        progress = (stats_dict['completed_chunks'] / stats_dict['total_chunks']) * 100
        expected = (i + 1) * 10  # 10%, 20%, 30%, 40%, 50%
        print(f"  Chunk {i+1}/5 completed: {progress:.1f}% (should be {expected:.0f}%)")

    stats_dict = stats.to_dict()
    progress = (stats_dict['completed_chunks'] / stats_dict['total_chunks']) * 100
    print(f"\nAfter translation phase:")
    print(f"  Completed chunks: {stats_dict['completed_chunks']}/{stats_dict['total_chunks']}")
    print(f"  Progress: {progress:.1f}% (should be 50.0%)")

    if abs(progress - 50.0) > 0.1:
        print(f"  [ERROR] Expected 50%, got {progress:.1f}%")
        return False
    else:
        print(f"  [OK] Correct!")

    # Start phase 2 (refinement)
    print("\n--- PHASE 2: REFINEMENT ---")
    stats.refinement_phase = True
    stats.refinement_chunks_completed = 0

    stats_dict = stats.to_dict()
    progress = (stats_dict['completed_chunks'] / stats_dict['total_chunks']) * 100
    print(f"After switching to refinement phase:")
    print(f"  Progress: {progress:.1f}% (should still be 50.0%)")

    # Complete phase 2 (refinement)
    for i in range(5):
        stats.refinement_chunks_completed = i + 1
        stats_dict = stats.to_dict()
        progress = (stats_dict['completed_chunks'] / stats_dict['total_chunks']) * 100
        expected = 50.0 + ((i + 1) * 10.0)  # 60%, 70%, 80%, 90%, 100%
        print(f"  Chunk {i+1}/5 refined: {progress:.1f}% (should be {expected:.0f}%)")

    stats_dict = stats.to_dict()
    progress = (stats_dict['completed_chunks'] / stats_dict['total_chunks']) * 100
    print(f"\nFinal state:")
    print(f"  Completed chunks: {stats_dict['completed_chunks']}/{stats_dict['total_chunks']}")
    print(f"  Progress: {progress:.1f}% (should be 100.0%)")

    if abs(progress - 100.0) > 0.1:
        print(f"  [ERROR] Expected 100%, got {progress:.1f}%")
        return False
    else:
        print(f"  [OK] Correct!")

    print("\n" + "=" * 60)
    print("[SUCCESS] All tests passed!")
    print("=" * 60)
    return True


def test_single_phase_epub_progress():
    """Test that single-phase EPUB mode still works correctly."""

    print("\n" + "=" * 60)
    print("Testing EPUB TranslationMetrics WITHOUT Refinement")
    print("=" * 60)

    # Create metrics without refinement
    stats = TranslationMetrics()
    stats.total_chunks = 5
    stats.enable_refinement = False

    print(f"\nInitial state:")
    stats_dict = stats.to_dict()
    print(f"  Total chunks: {stats_dict['total_chunks']}")
    print(f"  Completed chunks: {stats_dict['completed_chunks']}")
    print(f"  Progress: {(stats_dict['completed_chunks'] / stats_dict['total_chunks']) * 100:.1f}%")

    # Complete all chunks
    for i in range(5):
        stats.successful_first_try = i + 1
        stats.record_processed()  # Mark chunk as processed
        stats_dict = stats.to_dict()
        progress = (stats_dict['completed_chunks'] / stats_dict['total_chunks']) * 100
        expected = (i + 1) * 20.0
        print(f"  Chunk {i+1}/5 completed: {progress:.1f}% (should be {expected:.0f}%)")

    stats_dict = stats.to_dict()
    progress = (stats_dict['completed_chunks'] / stats_dict['total_chunks']) * 100
    print(f"\nFinal state:")
    print(f"  Progress: {progress:.1f}% (should be 100.0%)")

    if abs(progress - 100.0) > 0.1:
        print(f"  [ERROR] Expected 100%, got {progress:.1f}%")
        return False
    else:
        print(f"  [OK] Correct!")

    print("\n" + "=" * 60)
    print("[SUCCESS] Single-phase test passed!")
    print("=" * 60)
    return True


if __name__ == "__main__":
    success1 = test_single_phase_epub_progress()
    success2 = test_epub_refinement_progress()

    if success1 and success2:
        print("\n[SUCCESS] All EPUB progress tracking tests passed!")
    else:
        print("\n[FAILED] Some tests failed!")
        exit(1)