test_translation.py

"""
Test script for translation functionality with Google TranslateGemma model
"""
from hf_api import HuggingFaceAPI
from utils import load_settings
import os
import sys

# Import Google Translate Gemma if available
try:
    from google_translate import GoogleTranslateGemma
    GOOGLE_TRANSLATE_AVAILABLE = True
except ImportError as e:
    print(f"Warning: Google Translate Gemma not available: {str(e)}")
    GOOGLE_TRANSLATE_AVAILABLE = False

# Test translation with Google TranslateGemma
def test_translategemma(text, source_lang, target_lang):
    """Test translation using Google Translate Gemma model"""
    if not GOOGLE_TRANSLATE_AVAILABLE:
        print("❌ Google Translate Gemma is not available")
        print("   Falling back to chat completion translation")
        return test_chat_completion_translation(text, source_lang, target_lang)
    
    try:
        print(f"🧪 Testing Google Translate Gemma")
        print("-" * 50)
        
        # Initialize the translator
        translator = GoogleTranslateGemma()
        
        print(f"\n📝 Translating from {source_lang} to {target_lang}")
        print(f"Original: {text}")
        
        # Perform translation
        translation = translator.translate_text(
            text=text,
            source_lang=source_lang,
            target_lang=target_lang
        )
        
        print(f"✅ Translation: {translation}")
        
    except Exception as e:
        print(f"❌ Google Translate Gemma failed: {str(e)}")
        print("   Falling back to chat completion translation")
        return test_chat_completion_translation(text, source_lang, target_lang)
    
    print("\n" + "=" * 50)
    print("🎉 Google Translate Gemma testing complete!")

def test_chat_completion_translation(text, source_lang, target_lang):
    """Test translation using chat completion fallback"""
    # Load API token
    settings_dir = os.path.join(os.path.dirname(__file__), 'settings')
    models_settings_file = os.path.join(settings_dir, 'models.json')
    settings = load_settings(models_settings_file)
    token = settings.get('huggingfaceToken')
    
    if not token:
        print("❌ No HuggingFace token found. Please set your token first.")
        print("   You can set it in the app's Settings tab")
        return
    
    # Initialize API
    api = HuggingFaceAPI(token=token)
    
    # Test models in order of preference
    models_to_test = [
        "google/translategemma-12b-it",
        "meta-llama/Llama-3.2-3B-Instruct",
        "microsoft/Phi-3-mini-4k-instruct",
        "google/gemma-2-2b-it"
    ]
    
    print(f"🧪 Testing translation with chat completion")
    print("-" * 50)
    
    for model_id in models_to_test:
        print(f"\n📝 Testing with model: {model_id}")
        print(f"Original: {text}")
        
        try:
            # Use the same translation logic as in app.py
            if "translategemma" in model_id.lower() and not GOOGLE_TRANSLATE_AVAILABLE:
                print("   ⚠️  Google Translate Gemma not available, skipping...")
                continue
            
            # Dynamic system prompt based on target and source language
            source_info = f" from {source_lang}" if source_lang != "Auto-detect" else ""
            system_prompt = f"You are a professional translator specializing in translating{source_info} to {target_lang}. Translate the given text accurately while preserving the original meaning and tone. Only provide the translation without any additional explanations."
            prompt = f"Translate the following text{source_info} to {target_lang}: {text}"
            
            messages = [
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ]
            
            response = api.chat_completion(
                model=model_id,
                messages=messages,
                max_tokens=1024,
                temperature=0.3
            )
            
            translation = response["choices"][0]["message"]["content"].strip()
            print(f"✅ Translation: {translation}")
            print(f"   ✅ Success with {model_id}!")
            return translation  # Return first successful translation
            
        except Exception as e:
            print(f"   ❌ Error with {model_id}: {str(e)}")
            continue
    
    print("\n❌ All models failed. Please check your token and model availability.")
    return None

def test_multiple_translations():
    """Test multiple translation scenarios"""
    print("\n🌍 Testing Multiple Translation Scenarios")
    print("=" * 60)
    
    test_cases = [
        {
            "text": "Hello, how are you today?",
            "source": "English",
            "target": "Spanish",
            "description": "English to Spanish"
        },
        {
            "text": "V nejhorším případě i k prasknutí čočky.",
            "source": "Czech",
            "target": "German",
            "description": "Czech to German"
        },
        {
            "text": "Bonjour, comment allez-vous?",
            "source": "French",
            "target": "English",
            "description": "French to English"
        },
        {
            "text": "这是一个测试。",
            "source": "Chinese (Simplified)",
            "target": "English",
            "description": "Chinese to English"
        },
        {
            "text": "¡Hola! ¿Cómo estás?",
            "source": "Spanish",
            "target": "Japanese",
            "description": "Spanish to Japanese"
        }
    ]
    
    results = []
    
    for i, case in enumerate(test_cases, 1):
        print(f"\n📝 Test {i}: {case['description']}")
        print(f"   Source ({case['source']}): {case['text']}")
        
        # Map language names to codes
        lang_code_map = {
            "English": "en",
            "Spanish": "es",
            "French": "fr",
            "German": "de-DE",
            "Chinese (Simplified)": "zh-CN",
            "Chinese (Traditional)": "zh-TW",
            "Japanese": "ja",
            "Korean": "ko",
            "Italian": "it",
            "Portuguese": "pt",
            "Russian": "ru",
            "Arabic": "ar",
            "Hindi": "hi",
            "Dutch": "nl",
            "Turkish": "tr",
            "Polish": "pl",
            "Vietnamese": "vi",
            "Thai": "th",
            "Indonesian": "id",
            "Greek": "el",
            "Hebrew": "he",
            "Czech": "cs",
            "Swedish": "sv",
            "Danish": "da",
            "Norwegian": "no",
            "Finnish": "fi"
        }
        
        source_code = lang_code_map.get(case['source'], 'en')
        target_code = lang_code_map.get(case['target'], 'en')
        
        translation = test_translategemma(
            text=case['text'],
            source_lang=source_code,
            target_lang=target_code
        )
        
        if translation:
            print(f"   Target ({case['target']}): {translation}")
            results.append({
                'case': case['description'],
                'original': case['text'],
                'translation': translation,
                'success': True
            })
        else:
            results.append({
                'case': case['description'],
                'original': case['text'],
                'translation': None,
                'success': False
            })
    
    # Summary
    print("\n" + "=" * 60)
    print("📊 Test Summary")
    print("-" * 60)
    
    successful = sum(1 for r in results if r['success'])
    total = len(results)
    
    print(f"Total tests: {total}")
    print(f"Successful: {successful}")
    print(f"Failed: {total - successful}")
    print(f"Success rate: {successful/total*100:.1f}%")
    
    if successful < total:
        print("\n❌ Some tests failed. Check your HuggingFace token and model availability.")
    else:
        print("\n✅ All tests passed successfully!")
    
    return results

if __name__ == "__main__":
    import sys
    
    print("🌐 Translation Test Suite")
    print("=" * 60)
    print()
    
    # Check if command line arguments were provided
    if len(sys.argv) > 1:
        # Run single test with provided arguments
        if len(sys.argv) >= 4:
            text = sys.argv[1]
            source_lang = sys.argv[2]
            target_lang = sys.argv[3]
            print(f"Running single test:")
            print(f"  Text: {text}")
            print(f"  Source: {source_lang}")
            print(f"  Target: {target_lang}")
            print()
            test_translategemma(text, source_lang, target_lang)
        else:
            print("Usage: python test_translation.py <text> <source_lang> <target_lang>")
            print("Example: python test_translation.py \"Hello world\" en es")
    else:
        # Run comprehensive test suite
        print("Running comprehensive translation tests...")
        print()
        
        # First, test a simple case
        print("\n" + "=" * 60)
        print("🧪 Quick Test")
        test_translategemma(
            text="Hello, world!",
            source_lang="en",
            target_lang="es"
        )
        
        # Then run multiple tests
        test_multiple_translations()
        
        print("\n" + "=" * 60)
        print("🎉 All tests completed!")
        print()
        print("To test a specific translation:")
        print("  python test_translation.py \"Your text here\" source_lang target_lang")
        print()
        print("To test in the app:")
        print("  1. Run: python app.py")
        print("  2. Go to Translation tab")
        print("  3. Set your HuggingFace token in Settings")
        print("  4. Try translating text")