Upload /tmp/hugging-demos-build-model_17slever17_translate-gemma-4-sub-e4b-uly2snze/app.py with huggingface_hub

tmp/hugging-demos-build-model_17slever17_translate-gemma-4-sub-e4b-uly2snze/app.py

@@ -0,0 +1,186 @@
+import spaces  # MUST come before any torch/CUDA-touching import
+import torch
+import gradio as gr
+
+from transformers import AutoModelForImageTextToText, AutoTokenizer
+
+MODEL_ID = "17slever17/translate-gemma-4-sub-e4b"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForImageTextToText.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+model.eval()
+
+SUPPORTED_LANGUAGES = [
+    "English", "Russian", "Spanish", "German", "Japanese",
+    "French", "Portuguese", "Chinese", "Dutch", "Italian", "Korean",
+]
+
+STYLES = ["neutral", "friendly", "official"]
+
+
+def _build_system_message(source_lang: str, target_lang: str, rules: str, style: str) -> str:
+    """Build the system message exactly as specified in the model card."""
+    parts = [f"TASK: Translate {source_lang} subtitles into {target_lang}."]
+    if rules.strip():
+        parts.append(f"RULES: {rules.strip()}")
+    parts.append(f"STYLE: {style}.")
+    parts.append(
+        "Translate only CURRENT_SOURCE. PREVIOUS_SOURCE and PREVIOUS_TRANSLATION are context only. "
+        "Preserve meaning, tone, slang, profanity, uncertainty, repetitions and incomplete speech. "
+        "Return only the final translation without labels or commentary."
+    )
+    return "\n".join(parts)
+
+
+def _build_user_message(prev_source: str, prev_translation: str, current_source: str) -> str:
+    """Build the user message following the model card's format."""
+    blocks = []
+    if prev_source.strip():
+        blocks.append(f"[PREVIOUS_SOURCE]\n{prev_source.strip()}")
+    if prev_translation.strip():
+        blocks.append(f"[PREVIOUS_TRANSLATION]\n{prev_translation.strip()}")
+    blocks.append(f"[CURRENT_SOURCE]\n{current_source.strip()}")
+    return "\n\n".join(blocks)
+
+
+@spaces.GPU(duration=60)
+def translate(
+    current_source: str,
+    source_language: str = "English",
+    target_language: str = "Russian",
+    previous_source: str = "",
+    previous_translation: str = "",
+    rules: str = "",
+    style: str = "friendly",
+) -> str:
+    """Translate a subtitle line using context-aware Gemma 4 Sub model.
+
+    Args:
+        current_source: The subtitle segment to translate.
+        source_language: Language of the source text.
+        target_language: Language to translate into.
+        previous_source: Previous source-language subtitles for context (optional).
+        previous_translation: Previous translated subtitles for context (optional).
+        rules: Additional rules (speaker name, gender, terminology, etc.) (optional).
+        style: Translation style — neutral, friendly, or official.
+    """
+    if not current_source.strip():
+        return "Please enter a subtitle line to translate."
+
+    system_msg = _build_system_message(source_language, target_language, rules, style)
+    user_msg = _build_user_message(previous_source, previous_translation, current_source)
+
+    messages = [
+        {"role": "system", "content": system_msg},
+        {"role": "user", "content": user_msg},
+    ]
+
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        return_tensors="pt",
+        return_dict=True,
+    ).to("cuda")
+
+    with torch.inference_mode():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=False,
+        )
+
+    generated_tokens = output[0, inputs["input_ids"].shape[1]:]
+    translation = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
+    return translation
+
+
+CSS = """
+#col-container { max-width: 900px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        "# 🌐 Translate Gemma 4 Sub — Context-Aware Subtitle Translation\n"
+        "A multilingual Gemma 4 fine-tune specialized for subtitle translation. "
+        "Provide previous source and translation context for natural, coherent translations."
+    )
+
+    with gr.Column(elem_id="col-container"):
+        with gr.Row():
+            source_language = gr.Dropdown(
+                label="Source Language",
+                choices=SUPPORTED_LANGUAGES,
+                value="English",
+            )
+            target_language = gr.Dropdown(
+                label="Target Language",
+                choices=SUPPORTED_LANGUAGES,
+                value="Russian",
+            )
+            style = gr.Dropdown(
+                label="Style",
+                choices=STYLES,
+                value="friendly",
+            )
+
+        current_source = gr.Textbox(
+            label="Current Source (to translate)",
+            placeholder="Yeah, well... I changed my mind.",
+            lines=2,
+        )
+
+        with gr.Accordion("Context & Rules (optional)", open=False):
+            previous_source = gr.Textbox(
+                label="Previous Source (context only)",
+                placeholder="I thought you said you weren't coming.",
+                lines=2,
+            )
+            previous_translation = gr.Textbox(
+                label="Previous Translation (context only)",
+                placeholder="Я думала, ты сказала, что не придёшь.",
+                lines=2,
+            )
+            rules = gr.Textbox(
+                label="Rules (speaker name, gender, terminology, etc.)",
+                placeholder="speaker gender: female",
+                lines=2,
+            )
+
+        run_btn = gr.Button("Translate", variant="primary")
+        output = gr.Textbox(
+            label="Translation",
+            lines=3,
+            interactive=False,
+        )
+
+        run_btn.click(
+            fn=translate,
+            inputs=[
+                current_source, source_language, target_language,
+                previous_source, previous_translation, rules, style,
+            ],
+            outputs=output,
+            api_name="translate",
+        )
+
+    gr.Examples(
+        examples=[
+            ["Yeah, well... I changed my mind.", "English", "Russian", "friendly"],
+            ["I thought you said you weren't coming.", "English", "Spanish", "neutral"],
+            ["Wait, hold on a second...", "English", "Japanese", "friendly"],
+            ["That's not what I meant at all.", "English", "German", "official"],
+            ["Could you repeat that, please?", "English", "French", "official"],
+        ],
+        inputs=[current_source, source_language, target_language, style],
+        outputs=output,
+        fn=translate,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+
+if __name__ == "__main__":
+    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)