xTHExBEASTx/gemma-4-e4b-it

1import os2from collections.abc import Iterator3from threading import Thread4 5import gradio as gr6import spaces7import torch8from transformers import AutoModelForMultimodalLM, AutoProcessor, BatchFeature9from transformers.generation.streamers import TextIteratorStreamer10 11MODEL_ID = "google/gemma-4-e4b-it"12 13processor = AutoProcessor.from_pretrained(MODEL_ID, use_fast=False)14model = AutoModelForMultimodalLM.from_pretrained(MODEL_ID, device_map="auto", dtype=torch.bfloat16)15 16IMAGE_FILE_TYPES = (".jpg", ".jpeg", ".png", ".webp")17AUDIO_FILE_TYPES = (".wav", ".mp3", ".flac", ".ogg")18VIDEO_FILE_TYPES = (".mp4", ".mov", ".avi", ".webm")19MAX_INPUT_TOKENS = int(os.getenv("MAX_INPUT_TOKENS", "10_000"))20 21THINKING_START = "<|channel>"22THINKING_END = "<channel|>"23 24# Special tokens to strip from decoded output (keeping thinking delimiters25# so that Gradio's reasoning_tags can find them on the frontend).26_KEEP_TOKENS = {THINKING_START, THINKING_END}27_STRIP_TOKENS = sorted(28    (t for t in processor.tokenizer.all_special_tokens if t not in _KEEP_TOKENS),29    key=len,30    reverse=True,  # longest first to avoid partial matches31)32 33 34def _strip_special_tokens(text: str) -> str:35    for tok in _STRIP_TOKENS:36        text = text.replace(tok, "")37    return text38 39 40def _classify_file(path: str) -> str | None:41    """Return media type string for a file path, or None if unsupported."""42    lower = path.lower()43    if lower.endswith(IMAGE_FILE_TYPES):44        return "image"45    if lower.endswith(AUDIO_FILE_TYPES):46        return "audio"47    if lower.endswith(VIDEO_FILE_TYPES):48        return "video"49    return None50 51 52def process_new_user_message(message: dict) -> list[dict]:53    """Build content list from the new user message with URL-based media references."""54    content: list[dict] = []55    for path in message.get("files", []):56        kind = _classify_file(path)57        if kind:58            content.append({"type": kind, "url": path})59    content.append({"type": "text", "text": message.get("text", "")})60    return content61 62 63def process_history(history: list[dict]) -> list[dict]:64    """Walk Gradio 6 history and build message list with URL-based media references."""65    messages: list[dict] = []66 67    for item in history:68        if item["role"] == "assistant":69            text_parts = [p["text"] for p in item["content"] if p.get("type") == "text"]70            messages.append(71                {72                    "role": "assistant",73                    "content": [{"type": "text", "text": " ".join(text_parts)}],74                }75            )76        else:77            user_content: list[dict] = []78            for part in item["content"]:79                if part.get("type") == "text":80                    user_content.append({"type": "text", "text": part["text"]})81                elif part.get("type") == "file":82                    filepath = part["file"]["path"]83                    kind = _classify_file(filepath)84                    if kind:85                        user_content.append({"type": kind, "url": filepath})86            if user_content:87                messages.append({"role": "user", "content": user_content})88 89    return messages90 91 92@spaces.GPU(duration=120)93@torch.inference_mode()94def _generate_on_gpu(inputs: BatchFeature, max_new_tokens: int, thinking: bool) -> Iterator[str]:95    inputs = inputs.to(device=model.device, dtype=torch.bfloat16)96 97    streamer = TextIteratorStreamer(98        processor,99        timeout=30.0,100        skip_prompt=True,101        skip_special_tokens=not thinking,102    )103    generate_kwargs = {104        **inputs,105        "streamer": streamer,106        "max_new_tokens": max_new_tokens,107        "disable_compile": True,108    }109 110    exception_holder: list[Exception] = []111 112    def _generate() -> None:113        try:114            model.generate(**generate_kwargs)115        except Exception as e:  # noqa: BLE001116            exception_holder.append(e)117 118    thread = Thread(target=_generate)119    thread.start()120 121    chunks: list[str] = []122    for text in streamer:123        chunks.append(text)124        accumulated = "".join(chunks)125        if thinking:126            yield _strip_special_tokens(accumulated)127        else:128            yield accumulated129 130    thread.join()131    if exception_holder:132        msg = f"Generation failed: {exception_holder[0]}"133        raise gr.Error(msg)134 135 136def validate_input(message: dict) -> dict:137    has_text = bool(message.get("text", "").strip())138    has_files = bool(message.get("files"))139    if not (has_text or has_files):140        return gr.validate(False, "Please enter a message or upload a file.")141 142    files = message.get("files", [])143    kinds = [_classify_file(f) for f in files]144    kinds = [k for k in kinds if k is not None]145    unique_kinds = set(kinds)146 147    if len(unique_kinds) > 1:148        return gr.validate(False, "Please upload only one type of media (images, audio, or video) at a time.")149    if kinds.count("audio") > 1:150        return gr.validate(False, "Only one audio file can be uploaded at a time.")151    if kinds.count("video") > 1:152        return gr.validate(False, "Only one video file can be uploaded at a time.")153 154    return gr.validate(True, "")155 156 157def _has_media_type(messages: list[dict], media_type: str) -> bool:158    """Check if any message contains a content entry of the given media type."""159    return any(160        c.get("type") == media_type for m in messages for c in (m["content"] if isinstance(m["content"], list) else [])161    )162 163 164def generate(165    message: dict,166    history: list[dict],167    thinking: bool = False,168    max_new_tokens: int = 1024,169    max_soft_tokens: int = 280,170    system_prompt: str = "",171) -> Iterator[str]:172 173    messages: list[dict] = []174    if system_prompt:175        messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]})176 177    messages.extend(process_history(history))178    messages.append({"role": "user", "content": process_new_user_message(message)})179 180    template_kwargs: dict = {181        "tokenize": True,182        "return_dict": True,183        "return_tensors": "pt",184        "add_generation_prompt": True,185        "load_audio_from_video": _has_media_type(messages, "video"),186        "processor_kwargs": {"images_kwargs": {"max_soft_tokens": max_soft_tokens}},187    }188    if thinking:189        template_kwargs["enable_thinking"] = True190 191    inputs = processor.apply_chat_template(messages, **template_kwargs)192 193    n_tokens = inputs["input_ids"].shape[1]194    if n_tokens > MAX_INPUT_TOKENS:195        msg = f"Input too long ({n_tokens} tokens). Maximum is {MAX_INPUT_TOKENS} tokens."196        raise gr.Error(msg)197 198    yield from _generate_on_gpu(inputs=inputs, max_new_tokens=max_new_tokens, thinking=thinking)199 200 201examples = [202    # --- Text-only examples ---203    [204        {205            "text": "What is the capital of France?",206            "files": [],207        }208    ],209    [210        {211            "text": "What is the water formula?",212            "files": [],213        }214    ],215    [216        {217            "text": "Explain quantum entanglement in simple terms.",218            "files": [],219        }220    ],221    [222        {223            "text": "I want to do a car wash that is 50 meters away, should I walk or drive?",224            "files": [],225        }226    ],227    [228        {229            "text": "Write a poem about beer with 4 stanzas. Format the title as an H2 markdown heading and bold the first line of each stanza.",230            "files": [],231        }232    ],233    # --- Single-image examples ---234    [235        {236            "text": "Describe this image.",237            "files": ["https://news.bbc.co.uk/media/images/38107000/jpg/_38107299_ronaldogoal_ap_300.jpg"],238        }239    ],240    [241        {242            "text": "What is the city in this image? Describe what you see.",243            "files": ["https://imgmd.net/images/v1/guia/1698673/rio-de-janeiro-4-c.jpg"],244        }245    ],246    # --- Multi-image examples ---247    [248        {249            "text": "What are the key similarities between these three images?",250            "files": [251                "https://news.bbc.co.uk/media/images/38107000/jpg/_38107299_ronaldogoal_ap_300.jpg",252                "https://ogimg.infoglobo.com.br/in/12547538-502-0e0/FT1086A/94-8705-14.jpg",253                "https://amazonasatual.com.br/wp-content/uploads/2021/01/Pele.jpg",254            ],255        }256    ],257    # --- Audio examples ---258    [259        {260            "text": "Transcribe the audio.",261            "files": [262                "https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/bcn_weather.mp3"263            ],264        }265    ],266    [267        {268            "text": "Translate to Dutch.",269            "files": [270                "https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/bcn_weather.mp3"271            ],272        }273    ],274    # --- Video examples ---275    [276        {277            "text": "What is happening in this video?",278            "files": ["https://huggingface.co/datasets/merve/vlm_test_images/resolve/main/concert.mp4"],279        }280    ],281]282 283demo = gr.ChatInterface(284    fn=generate,285    validator=validate_input,286    chatbot=gr.Chatbot(287        scale=1,288        latex_delimiters=[289            {"left": "$$", "right": "$$", "display": True},290            {"left": "$", "right": "$", "display": False},291            {"left": "\\(", "right": "\\)", "display": False},292            {"left": "\\[", "right": "\\]", "display": True},293        ],294        reasoning_tags=[(THINKING_START, THINKING_END)],295    ),296    textbox=gr.MultimodalTextbox(297        sources=["upload", "microphone"],298        file_types=[*IMAGE_FILE_TYPES, *AUDIO_FILE_TYPES, *VIDEO_FILE_TYPES],299        file_count="multiple",300        autofocus=True,301    ),302    multimodal=True,303    additional_inputs=[304        gr.Checkbox(label="Thinking", value=False),305        gr.Slider(label="Max New Tokens", minimum=100, maximum=4000, step=10, value=2000),306        gr.Dropdown(307            label="Image Token Budget",308            info="Higher values preserve more visual detail (useful for OCR/documents). Lower values are faster.",309            choices=[70, 140, 280, 560, 1120],310            value=280,311        ),312        gr.Textbox(label="System Prompt", value=""),313    ],314    additional_inputs_accordion=gr.Accordion("Settings", open=True),315    stop_btn=False,316    title="Gemma 4 E4B It",317    examples=examples,318    run_examples_on_click=False,319    cache_examples=False,320    delete_cache=(1800, 1800),321)322 323if __name__ == "__main__":324    demo.launch(css_paths="style.css", max_file_size="20mb")325