#!/usr/bin/env python import os import gradio as gr import spaces import torch from transformers import AutoModelForImageTextToText, AutoProcessor MAX_NEW_TOKENS_LIMIT = int(os.getenv("MAX_NEW_TOKENS_LIMIT", "2000")) MAX_NEW_TOKENS_DEFAULT = int(os.getenv("MAX_NEW_TOKENS_DEFAULT", "400")) DEFAULT_SOURCE_LANG = os.getenv("DEFAULT_SOURCE_LANG", "en") DEFAULT_TARGET_LANG = os.getenv("DEFAULT_TARGET_LANG", "fr_FR") ALLOWED_MODEL_IDS = { "google/translategemma-4b-it", "google/translategemma-12b-it", "google/translategemma-27b-it", } model_id = os.getenv("MODEL_ID", "google/translategemma-27b-it") if model_id not in ALLOWED_MODEL_IDS: msg = f"MODEL_ID={model_id!r} is not supported. Choose one of: {', '.join(sorted(ALLOWED_MODEL_IDS))}" raise ValueError(msg) processor = AutoProcessor.from_pretrained(model_id) model = AutoModelForImageTextToText.from_pretrained(model_id, device_map="auto") LANG_CODE_TO_NAME = { "en": "English", "ar_EG": "Arabic (Egypt)", "ar_SA": "Arabic (Saudi Arabia)", "bg_BG": "Bulgarian", "bn_IN": "Bengali (India)", "ca_ES": "Catalan (Spain)", "cs_CZ": "Czech (Czechia)", "da_DK": "Danish (Denmark)", "de_DE": "German (Germany)", "el_GR": "Greek (Greece)", "es_MX": "Spanish (Mexico)", "et_EE": "Estonian (Estonia)", "fa_IR": "Persian (Iran)", "fi_FI": "Finnish (Finland)", "fil_PH": "Filipino (Philippines)", "fr_CA": "French (Canada)", "fr_FR": "French (France)", "gu_IN": "Gujarati (India)", "he_IL": "Hebrew (Israel)", "hi_IN": "Hindi (India)", "hr_HR": "Croatian (Croatia)", "hu_HU": "Hungarian (Hungary)", "id_ID": "Indonesian (Indonesia)", "is_IS": "Icelandic (Iceland)", "it_IT": "Italian (Italy)", "ja_JP": "Japanese (Japan)", "kn_IN": "Kannada (India)", "ko_KR": "Korean (Korea)", "lt_LT": "Lithuanian (Lithuania)", "lv_LV": "Latvian (Latvia)", "ml_IN": "Malayalam (India)", "mr_IN": "Marathi (India)", "nl_NL": "Dutch (Netherlands)", "no_NO": "Norwegian (Norway)", "pa_IN": "Punjabi (India)", "pl_PL": "Polish (Poland)", "pt_BR": "Portuguese (Brazil)", "pt_PT": "Portuguese (Portugal)", "ro_RO": "Romanian (Romania)", "ru_RU": "Russian (Russia)", "sk_SK": "Slovak (Slovakia)", "sl_SI": "Slovenian (Slovenia)", "sr_RS": "Serbian (Serbia)", "sv_SE": "Swedish (Sweden)", "sw_KE": "Swahili (Kenya)", "sw_TZ": "Swahili (Tanzania)", "ta_IN": "Tamil (India)", "te_IN": "Telugu (India)", "th_TH": "Thai (Thailand)", "tr_TR": "Turkish (Turkey)", "uk_UA": "Ukrainian (Ukraine)", "ur_PK": "Urdu (Pakistan)", "vi_VN": "Vietnamese (Vietnam)", "zh_CN": "Chinese (Simplified, China)", "zh_TW": "Chinese (Traditional, Taiwan)", "zu_ZA": "Zulu (South Africa)", } LANG_CHOICES = [ ("English (en)", "en"), *sorted( [(f"{name} ({code})", code) for code, name in LANG_CODE_TO_NAME.items() if code != "en"], key=lambda x: x[0].lower(), ), ] _valid_lang_codes = set(LANG_CODE_TO_NAME.keys()) if DEFAULT_SOURCE_LANG not in _valid_lang_codes: msg = f"DEFAULT_SOURCE_LANG={DEFAULT_SOURCE_LANG!r} is not a supported language code." raise ValueError(msg) if DEFAULT_TARGET_LANG not in _valid_lang_codes: msg = f"DEFAULT_TARGET_LANG={DEFAULT_TARGET_LANG!r} is not a supported language code." raise ValueError(msg) def _build_messages(text: str, source_lang_code: str, target_lang_code: str) -> list[dict]: # TODO:Remove this once the chat template is fixed. # noqa: FIX002, TD002, TD003, TD007 # Temporary workaround: Model expects zh_CH instead of zh_CN due to a bug if source_lang_code == "zh_CN": source_lang_code = "zh_CH" if target_lang_code == "zh_CN": target_lang_code = "zh_CH" return [ { "role": "user", "content": [ { "type": "text", "source_lang_code": source_lang_code, "target_lang_code": target_lang_code, "text": text, } ], } ] def swap_languages(source: str, target: str) -> tuple[str, str]: return target, source def count_tokens(text: str, source_lang_code: str, target_lang_code: str) -> str: """Count input tokens without GPU. Returns a short info string.""" if not text: return "" if source_lang_code not in LANG_CODE_TO_NAME or target_lang_code not in LANG_CODE_TO_NAME: return "" messages = _build_messages(text, source_lang_code, target_lang_code) inputs = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_dict=True) return f"Input tokens: {len(inputs['input_ids'][0])}" @spaces.GPU(duration=120) @torch.inference_mode() def translate(text: str, source_lang_code: str, target_lang_code: str, max_new_tokens: int) -> str: if not text: raise gr.Error("Please enter text to translate") if source_lang_code not in LANG_CODE_TO_NAME: error_message = f"Invalid source language: {source_lang_code}" raise gr.Error(error_message) if target_lang_code not in LANG_CODE_TO_NAME: error_message = f"Invalid target language: {target_lang_code}" raise gr.Error(error_message) messages = _build_messages(text, source_lang_code, target_lang_code) 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_len = len(inputs["input_ids"][0]) if input_len + max_new_tokens > MAX_NEW_TOKENS_LIMIT: error_message = f"Input ({input_len} tokens) + max output ({max_new_tokens} tokens) exceeds the total limit of {MAX_NEW_TOKENS_LIMIT} tokens." raise gr.Error(error_message) generation = model.generate( **inputs, do_sample=False, max_new_tokens=max_new_tokens, cache_implementation="dynamic" ) generation = generation[0][input_len:] return processor.decode(generation, skip_special_tokens=True) with gr.Blocks() as demo: gr.Markdown("# Translategemma-27b-it") with gr.Row(): source_lang_code = gr.Dropdown(label="Source Language", choices=LANG_CHOICES, value=DEFAULT_SOURCE_LANG) swap_button = gr.Button("⇆", elem_id="swap-btn") target_lang_code = gr.Dropdown(label="Target Language", choices=LANG_CHOICES, value=DEFAULT_TARGET_LANG) max_new_tokens = gr.Slider( label="Max New Tokens", info="Higher values allow longer translations but take more time", minimum=50, maximum=MAX_NEW_TOKENS_LIMIT, step=10, value=MAX_NEW_TOKENS_DEFAULT, ) with gr.Row(): with gr.Column(): text = gr.Textbox(label="Input", lines=10, placeholder="Enter text to translate") token_info = gr.Textbox(label="Token Count", lines=1) translate_button = gr.Button("Translate", variant="primary") with gr.Column(): output = gr.Textbox(label="Translation", lines=10, placeholder="Translation will appear here") token_count_inputs = [text, source_lang_code, target_lang_code] for component in token_count_inputs: component.change(fn=count_tokens, inputs=token_count_inputs, outputs=token_info) swap_button.click( fn=swap_languages, inputs=[source_lang_code, target_lang_code], outputs=[source_lang_code, target_lang_code], ) translate_button.click( fn=translate, inputs=[text, source_lang_code, target_lang_code, max_new_tokens], outputs=output, ) if __name__ == "__main__": demo.launch(css_paths="style.css")