""" GemmaTranslate Service - Test Only Version This version excludes web endpoints to avoid FastAPI dependency during testing """ import modal import re from typing import Dict, Tuple import logging import time # Setup logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Create Modal app app = modal.App("translate-gemma-test") # Create persistent volume for model storage model_volume = modal.Volume.from_name("gemmatranslate-models", create_if_missing=True) # HuggingFace token for model download HF_TOKEN = "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" # Download model during image build to avoid runtime network access issues def download_model(): from transformers import AutoProcessor, AutoModelForImageTextToText import os import torch model_name = "google/translategemma-4b-it" cache_dir = "/model-cache" print(f"Downloading {model_name} to image...") os.makedirs(cache_dir, exist_ok=True) processor = AutoProcessor.from_pretrained(model_name, token=HF_TOKEN, cache_dir=cache_dir) model = AutoModelForImageTextToText.from_pretrained( model_name, token=HF_TOKEN, cache_dir=cache_dir, torch_dtype=torch.bfloat16 ) print("Model downloaded successfully!") return cache_dir # Docker image with ML dependencies and pre-downloaded model image = ( modal.Image.debian_slim(python_version="3.11") .pip_install( "transformers>=4.46.0", # Upgraded for Gemma tokenizer compatibility "torch>=2.6.0", # Required for Gemma3 masking functions "sentencepiece==0.2.0", "protobuf==5.28.2", "accelerate>=0.34.2", "safetensors>=0.4.3", "pillow", # Required for AutoImageProcessor ) .run_function(download_model, secrets=[modal.Secret.from_dict({"HUGGING_FACE_HUB_TOKEN": HF_TOKEN})]) ) # Model cache path in the volume MODEL_CACHE_PATH = "/models" class HTMLTagPreserver: """Utility class to preserve HTML tags during translation""" @staticmethod def extract_tags(content: str) -> Tuple[str, Dict[str, str]]: """Extract HTML tags and replace with placeholders""" tag_pattern = r'<[^>]+>' tags = {} tag_counter = 0 def replace_tag(match): nonlocal tag_counter tag = match.group(0) placeholder = f"__TAG_{tag_counter}__" tags[placeholder] = tag tag_counter += 1 return placeholder clean_content = re.sub(tag_pattern, replace_tag, content) return clean_content, tags @staticmethod def restore_tags(translated_content: str, tags: Dict[str, str]) -> str: """Restore HTML tags in translated content""" result = translated_content for placeholder, tag in tags.items(): result = result.replace(placeholder, tag) return result @app.function( image=image, gpu="T4", timeout=180, # 3 minutes for model load + generation memory=20480, volumes={MODEL_CACHE_PATH: model_volume}, ) def translate_gemma( content: str, source_lang: str, target_lang: str, tone: str = "neutral" ) -> Dict[str, any]: """ Translate using Google's TranslateGemma 4B instruction-tuned model """ from transformers import AutoProcessor, AutoModelForImageTextToText import os import torch start_time = time.time() try: logger.info(f"Translating from {source_lang} to {target_lang} using TranslateGemma 4B") # Extract HTML tags clean_content, tags = HTMLTagPreserver.extract_tags(content) if not clean_content.strip(): logger.warning("Empty content after tag extraction") return { "translation": content, "tokens_used": 0, "processing_time_ms": 0, "model": "gemmatranslate" } # TranslateGemma 4B - load from image cache (pre-downloaded during build) model_name = "google/translategemma-4b-it" image_cache_dir = "/model-cache" logger.info("Loading model from image cache...") load_start = time.time() # Load processor and model from cache processor = AutoProcessor.from_pretrained( model_name, cache_dir=image_cache_dir, local_files_only=True ) model = AutoModelForImageTextToText.from_pretrained( model_name, cache_dir=image_cache_dir, local_files_only=True, torch_dtype=torch.bfloat16, device_map="auto" ) logger.info(f"Model loaded in {time.time() - load_start:.2f}s") # Prepare translation using TranslateGemma chat template messages = [ { "role": "user", "content": [ { "type": "text", "source_lang_code": source_lang, "target_lang_code": target_lang, "text": clean_content } ] } ] # Apply chat template inputs = processor.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt" ).to(model.device, dtype=torch.bfloat16) input_tokens = inputs['input_ids'].shape[1] logger.info(f"Input tokens: {input_tokens}") # Generate translation with torch.inference_mode(): output = model.generate( **inputs, max_new_tokens=512, do_sample=False, ) # Decode output - CRITICAL: Remove input tokens first (official HuggingFace approach) input_len = len(inputs['input_ids'][0]) generation = output[0][input_len:] translated = processor.decode(generation, skip_special_tokens=True) translated = translated.strip() # Restore HTML tags translated_with_tags = HTMLTagPreserver.restore_tags(translated, tags) output_tokens = len(generation) total_tokens = input_tokens + output_tokens processing_time = int((time.time() - start_time) * 1000) logger.info( f"Translation completed in {processing_time}ms using {total_tokens} tokens " f"(input: {input_tokens}, output: {output_tokens})" ) return { "translation": translated_with_tags, "tokens_used": total_tokens, "processing_time_ms": processing_time, "model": "gemmatranslate", "source_lang": source_lang, "target_lang": target_lang } except Exception as e: logger.error(f"Translation error: {str(e)}", exc_info=True) raise RuntimeError(f"Translation failed: {str(e)}") from e # Local entrypoint for testing @app.local_entrypoint() def test_translation(): """Test the translation service locally""" test_text = "The quick brown fox jumps over the lazy dog. This is a simple test sentence." print(f"\n{'='*70}") print(f" TranslateGemma 4B Translation Test") print(f"{'='*70}\n") print(f"📝 Input text:") print(f" {test_text}\n") print(f"🌐 Translation: en -> he\n") # Run translation print(f"🔄 Starting translation...") start = time.time() result = translate_gemma.remote(test_text, "en", "he", "neutral") elapsed = time.time() - start print(f"\n{'='*70}") print(f" Translation Results") print(f"{'='*70}\n") print(f"✅ Translation successful!\n") print(f"📝 Original:") print(f" {test_text}\n") print(f"🌐 Translated:") print(f" {result['translation']}\n") print(f"📊 Metrics:") print(f" Tokens used: {result['tokens_used']}") print(f" Processing time (server): {result['processing_time_ms']}ms") print(f" Total time (with network): {elapsed*1000:.0f}ms") print(f" Model: {result['model']}\n") print(f"{'='*70}") print(f" Test completed successfully! ✅") print(f"{'='*70}\n")