blank/test_translation.py

#!/usr/bin/env python3
"""Test translation with glossary terms."""

import sys

# Add torch path first
torch_path = '/root/.local/lib/python3.11/site-packages'
if torch_path not in sys.path:
    sys.path.insert(0, torch_path)

# Add src path
sys.path.insert(0, '/mnt/code/223-236-template-6/src')

from glossary.models import Glossary, GlossaryEntry, TermCategory
from translator.pipeline import TranslationPipeline
from translator.engine import TranslationEngine


def main():
    # Initialize components
    model_path = '/mnt/code/223-236-template-6/phase0-test/models/m2m100_418M'
    engine = TranslationEngine(model_path=model_path)
    glossary = Glossary()

    # Add terms
    glossary.add(GlossaryEntry(source='林风', target='Lin Feng', category=TermCategory.CHARACTER))
    glossary.add(GlossaryEntry(source='青云宗', target='Qingyun Sect', category=TermCategory.ORGANIZATION))
    glossary.add(GlossaryEntry(source='外门弟子', target='outer disciple', category=TermCategory.OTHER))

    print('=' * 50)
    print('术语表已添加以下术语:')
    for entry in glossary.get_all():
        print(f'  - {entry.source} -> {entry.target}')
    print('=' * 50)

    # Create pipeline
    pipeline = TranslationPipeline(
        engine=engine,
        glossary=glossary,
        src_lang='zh',
        tgt_lang='en'
    )

    # Test translation
    test_text = '林风是青云宗的一名外门弟子'
    print(f'\n输入: {test_text}')
    print('=' * 50)

    result = pipeline.translate(test_text, return_details=True)

    print(f'输出: {result.translated}')
    print(f'使用的术语: {result.terms_used if result.terms_used else []}')
    print('=' * 50)

    # Check for _en__ prefix issue
    if '_en__' in result.translated:
        print('\n❌ 问题: 发现 _en__ 前缀残留!')
        return 1
    else:
        print('\n✅ 通过: 没有 _en__ 前缀问题!')
        return 0


if __name__ == '__main__':
    sys.exit(main() or 0)