Spaces:

Bertug1911
/

Model-Interface-for_all_models

Sleeping

Bertug1911

Update app.py

1e37962 verified about 1 month ago

18.3 kB

# Import before torch/CUDA initialization on Hugging Face Spaces.
import spaces

import gc
import queue
import re
import threading
import traceback

import gradio as gr
import torch
from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForCausalLM,
    TextIteratorStreamer,
)

# ============================================================
# CONFIGURATION
# ============================================================

MODELS = {
    "BrtGPT-1-Pre": "Bertug1911/BrtGPT-1-Pre",
    "BrtGPT-1-Pre-Code": "Bertug1911/BrtGPT-1-Pre-Code",
    "BrtGPT-1-0719": "HuggingFaceTB/SmolLM2-1.7B-Instruct",
    "BrtGPT-1-0906": "Bertug1911/BrtGPT-1-0906",
    "BrtGPT-Conversation-1": "Bertug1911/BrtGPT-Conversation-1",
    "BrtGPT-Conversation-1.5": "Bertug1911/BrtGPT-Conversation-1.5",
    "BrtGPT-Conversation-2": "Bertug1911/BrtGPT-Conversation-2",
    "CodePiegon": "Bertug1911/CodePiegon",
}
DEFAULT_MODEL = "BrtGPT-1-Pre"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
REPETITION_PENALTY = 1.1

# Per-model maximums. Global minimums: top_p=0.1, top_k=1,
# temperature=0.1, max_new_tokens=1. Adjust these limits as desired.
MODEL_LIMITS = {
    "BrtGPT-1-Pre":             {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "BrtGPT-1-Pre-Code":        {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "BrtGPT-1-0719":            {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "BrtGPT-1-0906":            {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "BrtGPT-Conversation-1":    {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "BrtGPT-Conversation-1.5":  {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "BrtGPT-Conversation-2":    {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "CodePiegon":               {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 128},
}
DEFAULTS = {"top_p": 0.75, "top_k": 12, "temperature": 0.7, "max_new_tokens": 128}

_current_model_id = None
_current_model = None
_current_tokenizer = None
_model_lock = threading.Lock()


def get_limits(model_name):
    return MODEL_LIMITS.get(model_name, MODEL_LIMITS[DEFAULT_MODEL])


def update_model_controls(model_name):
    limits = get_limits(model_name)
    return (
        gr.update(minimum=0.1, maximum=limits["top_p_max"], value=min(DEFAULTS["top_p"], limits["top_p_max"])),
        gr.update(minimum=1, maximum=limits["top_k_max"], value=min(DEFAULTS["top_k"], limits["top_k_max"]), step=1),
        gr.update(minimum=0.1, maximum=limits["temperature_max"], value=min(DEFAULTS["temperature"], limits["temperature_max"])),
        gr.update(minimum=1, maximum=limits["max_new_tokens_max"], value=min(DEFAULTS["max_new_tokens"], limits["max_new_tokens_max"]), step=1),
    )


# ============================================================
# MODEL LOADING + TOKEN ID SAFETY CHECKS
# ============================================================

def load_model(model_name):
    global _current_model_id, _current_model, _current_tokenizer

    if model_name not in MODELS:
        raise ValueError(f"Unknown model: {model_name}")

    model_id = MODELS[model_name]
    if _current_model_id == model_id and _current_model is not None:
        return _current_tokenizer, _current_model

    # Release the previous model before loading another one.
    if _current_model is not None:
        del _current_model
    _current_model = None
    _current_tokenizer = None
    _current_model_id = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    print(f"Loading tokenizer: {model_id}")
    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
    config = AutoConfig.from_pretrained(model_id)

    # Important: if the tokenizer has more IDs than the model config,
    # set the config vocabulary size BEFORE loading weights. This avoids
    # the common 32000-vs-32001 embedding shape mismatch when the checkpoint
    # actually contains the extra embedding row.
    tokenizer_size = len(tokenizer)
    if tokenizer_size > int(config.vocab_size):
        print(f"Adjusting config vocab_size: {config.vocab_size} -> {tokenizer_size}")
        config.vocab_size = tokenizer_size

    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        config=config,
        torch_dtype=DTYPE,
        low_cpu_mem_usage=True,
    )

    embedding_count = model.get_input_embeddings().num_embeddings
    if tokenizer_size > embedding_count:
        # Defensive fallback for architectures whose embedding table was not
        # sized from config. New rows are initialized; see note in UI errors.
        model.resize_token_embeddings(tokenizer_size)
        embedding_count = model.get_input_embeddings().num_embeddings

    # Validate all tokenizer IDs before the first CUDA operation.
    vocab = tokenizer.get_vocab()
    invalid = [(token, idx) for token, idx in vocab.items() if idx < 0 or idx >= embedding_count]
    if invalid:
        del model
        gc.collect()
        raise ValueError(
            f"Tokenizer/model vocabulary mismatch for {model_name}: {len(invalid)} token IDs "
            f"are outside embedding range [0, {embedding_count - 1}]. "
            "Check that this model repository's tokenizer and weights come from the same revision."
        )

    if tokenizer.pad_token_id is None:
        if tokenizer.eos_token_id is not None:
            tokenizer.pad_token = tokenizer.eos_token
        else:
            tokenizer.add_special_tokens({"pad_token": "<|pad|>"})
            if len(tokenizer) > model.get_input_embeddings().num_embeddings:
                model.resize_token_embeddings(len(tokenizer))

    model.to(DEVICE)
    model.eval()

    _current_tokenizer = tokenizer
    _current_model = model
    _current_model_id = model_id
    print(
        f"Model ready: {model_id} on {DEVICE} | tokenizer={len(tokenizer)} | "
        f"embeddings={model.get_input_embeddings().num_embeddings}"
    )
    return tokenizer, model


# ============================================================
# OUTPUT FORMATTING
# ============================================================

CODE_BLOCK_PATTERN = re.compile(r"```([^\n`]*)\n(.*?)```", re.DOTALL)


def format_output(text):
    def replace_code_block(match):
        language = match.group(1).strip() or "text"
        code = match.group(2).rstrip("\n")
        longest_run = max((len(run) for run in re.findall(r"`+", code)), default=0)
        fence = "`" * max(3, longest_run + 1)
        return f"\n\n### `<{language.upper()}>`\n\n{fence}{language}\n{code}\n{fence}\n\n"

    return CODE_BLOCK_PATTERN.sub(replace_code_block, text).strip() or "The model returned an empty response."


# Preserve the original marker cleanup behavior, but only when those markers
# actually appear in the model output.
def normalize_token_markers(text):
    if "Ġ" in text or "Ċ" in text:
        text = text.replace(" ", "")
        text = text.replace("Ġ", " ").replace("Ċ", "\n")
    return text


# ============================================================
# GENERATION
# ============================================================

@spaces.GPU
def generate(prompt, model_name, top_p, top_k, temperature, max_new_tokens):
    if not prompt or not prompt.strip():
        yield "Please enter a message."
        return
    if model_name not in MODELS:
        yield "Please select a valid model."
        return

    limits = get_limits(model_name)
    # Clamp all UI values again server-side; browser inputs are not trusted.
    top_p = max(0.1, min(float(top_p), limits["top_p_max"]))
    top_k = max(1, min(int(top_k), limits["top_k_max"]))
    temperature = max(0.1, min(float(temperature), limits["temperature_max"]))
    max_new_tokens = max(1, min(int(max_new_tokens), limits["max_new_tokens_max"]))

    yield "Loading model and generating..."

    try:
        # Keep the lock for the complete operation; only one request/model
        # generation runs at once, avoiding unsafe concurrent CUDA use.
        with _model_lock:
            tokenizer, model = load_model(model_name)
            messages = [{"role": "user", "content": prompt.strip()}]

            if getattr(tokenizer, "chat_template", None):
                inputs = tokenizer.apply_chat_template(
                    messages,
                    tokenize=True,
                    add_generation_prompt=True,
                    return_tensors="pt",
                    return_dict=True,
                )
            else:
                inputs = tokenizer(prompt.strip(), return_tensors="pt")

            inputs = {key: value.to(DEVICE) for key, value in inputs.items()}
            input_ids = inputs.get("input_ids")
            if input_ids is None or input_ids.numel() == 0:
                raise ValueError("Tokenizer returned empty input_ids.")

            embedding_count = model.get_input_embeddings().num_embeddings
            min_id = int(input_ids.min().item())
            max_id = int(input_ids.max().item())
            if min_id < 0 or max_id >= embedding_count:
                raise ValueError(
                    f"Input token ID range [{min_id}, {max_id}] exceeds model embedding range "
                    f"[0, {embedding_count - 1}]. Tokenizer and model files are incompatible."
                )

            model_context = getattr(model.config, "max_position_embeddings", None)
            if model_context is None:
                model_context = getattr(model.config, "n_positions", None)
            if isinstance(model_context, int) and model_context > 1:
                # Keep the prompt inside the model context and leave room for generation.
                # This also prevents position-embedding index errors on very long prompts.
                max_input_tokens = model_context - 1
                if input_ids.shape[-1] > max_input_tokens:
                    for key in ("input_ids", "attention_mask", "token_type_ids"):
                        if key in inputs and inputs[key].ndim == 2:
                            inputs[key] = inputs[key][:, -max_input_tokens:]
                    input_ids = inputs["input_ids"]
                    print(
                        f"Prompt truncated to {max_input_tokens} tokens for model context "
                        f"limit {model_context}."
                    )
                max_new_tokens = min(max_new_tokens, max(1, model_context - input_ids.shape[-1]))

            pad_token_id = tokenizer.pad_token_id
            if pad_token_id is None:
                pad_token_id = tokenizer.eos_token_id

            generation_kwargs = {
                "max_new_tokens": max_new_tokens,
                "do_sample": True,
                "temperature": temperature,
                "top_p": top_p,
                "top_k": top_k,
                "repetition_penalty": REPETITION_PENALTY,
                "pad_token_id": pad_token_id,
                "use_cache": True,
            }
            if tokenizer.eos_token_id is not None:
                generation_kwargs["eos_token_id"] = tokenizer.eos_token_id

            streamer = TextIteratorStreamer(
                tokenizer,
                skip_prompt=True,
                skip_special_tokens=True,
                timeout=0.5,
            )
            worker_errors = []

            def generation_worker():
                try:
                    with torch.inference_mode():
                        model.generate(**inputs, streamer=streamer, **generation_kwargs)
                except Exception as exc:
                    worker_errors.append(exc)
                    try:
                        streamer.on_finalized_text("", stream_end=True)
                    except Exception:
                        pass

            worker = threading.Thread(target=generation_worker, daemon=True)
            worker.start()
            raw_response = ""

            while True:
                try:
                    chunk = next(streamer)
                except StopIteration:
                    break
                except queue.Empty:
                    if not worker.is_alive():
                        if worker_errors:
                            raise worker_errors[0]
                        break
                    continue

                if chunk:
                    raw_response += chunk
                    yield format_output(normalize_token_markers(raw_response))

            worker.join()
            if worker_errors:
                raise worker_errors[0]
            yield format_output(normalize_token_markers(raw_response))

    except torch.cuda.OutOfMemoryError:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        yield "GPU memory is insufficient. Select a smaller model or reduce Max New Tokens."
    except Exception as exc:
        print("Generation error:", repr(exc))
        traceback.print_exc()
        message = str(exc)
        if "device-side assert" in message.lower() or "cuda error" in message.lower():
            yield (
                "CUDA generation failed. The app checked token IDs before generation, so this may be "
                "a model/checkpoint incompatibility or a CUDA kernel issue. Restart the Space runtime "
                "and verify that the model weights and tokenizer are from the same repository revision."
            )
        else:
            yield f"Generation error ({type(exc).__name__}): {message}"


# ============================================================
# DARK THEME
# ============================================================

CSS = """
.gradio-container { max-width: 940px !important; margin: 0 auto !important; background: #0b0d12 !important; }
#app-header { text-align: center; padding: 22px 0 8px 0; }
#app-header h1 { font-size: 30px; font-weight: 750; letter-spacing: -1px; color: #f1f5f9; }
#app-subtitle { text-align: center; color: #929bad; margin-bottom: 22px; }
#model-selector, #settings-panel { border: 1px solid #292e3b; border-radius: 12px; background: #11141c; }
#input-area textarea { background: #11141c !important; border: 1px solid #292e3b !important; border-radius: 12px !important; color: #f1f5f9 !important; font-size: 15px !important; line-height: 1.7 !important; }
#output-area { border: 1px solid #292e3b; border-radius: 12px; background: #11141c; padding: 18px; }
#output-area h3 { color: #a5b4fc; font-size: 12px; font-weight: 700; letter-spacing: 1px; margin-bottom: 10px; }
#output-area pre { border: 1px solid #292e3b !important; border-radius: 10px !important; background: #090b10 !important; padding: 16px !important; overflow-x: auto !important; }
#output-area code { font-size: 13px !important; line-height: 1.7 !important; }
#generate-button { border-radius: 10px; min-height: 46px; font-weight: 650; }
footer { display: none !important; }
"""

with gr.Blocks(
    theme=gr.themes.Base(
        primary_hue="indigo",
        secondary_hue="slate",
        neutral_hue="slate",
    ).set(
        body_background_fill="#0b0d12",
        body_text_color="#e5e7eb",
        block_background_fill="#11141c",
        block_border_color="#292e3b",
        input_background_fill="#11141c",
        button_primary_background_fill="#6366f1",
        button_primary_background_fill_hover="#4f46e5",
    ),
    css=CSS,
    title="LLM Playground",
) as demo:
    gr.Markdown("# LLM Playground", elem_id="app-header")
    gr.Markdown("A simple, modern interface for language models.", elem_id="app-subtitle")

    model_selector = gr.Dropdown(
        choices=list(MODELS.keys()),
        value=DEFAULT_MODEL,
        label="MODEL",
        info="Choose which model to run.",
        elem_id="model-selector",
        interactive=True,
    )

    prompt_input = gr.Textbox(
        label="INPUT",
        placeholder="Ask a question, write code, or enter a prompt...",
        lines=5,
        max_lines=12,
        elem_id="input-area",
    )

    with gr.Accordion("Generation settings", open=True, elem_id="settings-panel"):
        with gr.Row():
            top_p_input = gr.Slider(
                minimum=0.1, maximum=get_limits(DEFAULT_MODEL)["top_p_max"], value=DEFAULTS["top_p"],
                step=0.05, label="Top P", info="Minimum 0.1; maximum depends on selected model."
            )
            top_k_input = gr.Slider(
                minimum=1, maximum=get_limits(DEFAULT_MODEL)["top_k_max"], value=DEFAULTS["top_k"],
                step=1, label="Top K", info="Minimum 1; maximum depends on selected model."
            )
        with gr.Row():
            temperature_input = gr.Slider(
                minimum=0.1, maximum=get_limits(DEFAULT_MODEL)["temperature_max"], value=DEFAULTS["temperature"],
                step=0.05, label="Temperature", info="Minimum 0.1; maximum depends on selected model."
            )
            max_new_tokens_input = gr.Slider(
                minimum=1, maximum=get_limits(DEFAULT_MODEL)["max_new_tokens_max"], value=DEFAULTS["max_new_tokens"],
                step=1, label="Max New Tokens", info="Minimum 1; maximum depends on selected model."
            )

    generate_button = gr.Button("Generate", variant="primary", elem_id="generate-button")
    output = gr.Markdown(value="Your model's response will appear here.", label="OUTPUT", elem_id="output-area")

    model_selector.change(
        fn=update_model_controls,
        inputs=[model_selector],
        outputs=[top_p_input, top_k_input, temperature_input, max_new_tokens_input],
    )

    generation_inputs = [
        prompt_input, model_selector, top_p_input, top_k_input,
        temperature_input, max_new_tokens_input,
    ]
    generate_button.click(fn=generate, inputs=generation_inputs, outputs=output)
    prompt_input.submit(fn=generate, inputs=generation_inputs, outputs=output)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        server_name="0.0.0.0",
        server_port=7860,
    )

	# Import before torch/CUDA initialization on Hugging Face Spaces.
	import spaces

	import gc
	import queue
	import re
	import threading
	import traceback

	import gradio as gr
	import torch
	from transformers import (
	AutoConfig,
	AutoTokenizer,
	AutoModelForCausalLM,
	TextIteratorStreamer,
	)

	# ============================================================
	# CONFIGURATION
	# ============================================================

	MODELS = {
	"BrtGPT-1-Pre": "Bertug1911/BrtGPT-1-Pre",
	"BrtGPT-1-Pre-Code": "Bertug1911/BrtGPT-1-Pre-Code",
	"BrtGPT-1-0719": "HuggingFaceTB/SmolLM2-1.7B-Instruct",
	"BrtGPT-1-0906": "Bertug1911/BrtGPT-1-0906",
	"BrtGPT-Conversation-1": "Bertug1911/BrtGPT-Conversation-1",
	"BrtGPT-Conversation-1.5": "Bertug1911/BrtGPT-Conversation-1.5",
	"BrtGPT-Conversation-2": "Bertug1911/BrtGPT-Conversation-2",
	"CodePiegon": "Bertug1911/CodePiegon",
	}
	DEFAULT_MODEL = "BrtGPT-1-Pre"
	DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
	DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
	REPETITION_PENALTY = 1.1

	# Per-model maximums. Global minimums: top_p=0.1, top_k=1,
	# temperature=0.1, max_new_tokens=1. Adjust these limits as desired.
	MODEL_LIMITS = {
	"BrtGPT-1-Pre": {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
	"BrtGPT-1-Pre-Code": {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
	"BrtGPT-1-0719": {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
	"BrtGPT-1-0906": {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
	"BrtGPT-Conversation-1": {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
	"BrtGPT-Conversation-1.5": {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
	"BrtGPT-Conversation-2": {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
	"CodePiegon": {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 128},
	}
	DEFAULTS = {"top_p": 0.75, "top_k": 12, "temperature": 0.7, "max_new_tokens": 128}

	_current_model_id = None
	_current_model = None
	_current_tokenizer = None
	_model_lock = threading.Lock()


	def get_limits(model_name):
	return MODEL_LIMITS.get(model_name, MODEL_LIMITS[DEFAULT_MODEL])


	def update_model_controls(model_name):
	limits = get_limits(model_name)
	return (
	gr.update(minimum=0.1, maximum=limits["top_p_max"], value=min(DEFAULTS["top_p"], limits["top_p_max"])),
	gr.update(minimum=1, maximum=limits["top_k_max"], value=min(DEFAULTS["top_k"], limits["top_k_max"]), step=1),
	gr.update(minimum=0.1, maximum=limits["temperature_max"], value=min(DEFAULTS["temperature"], limits["temperature_max"])),
	gr.update(minimum=1, maximum=limits["max_new_tokens_max"], value=min(DEFAULTS["max_new_tokens"], limits["max_new_tokens_max"]), step=1),
	)


	# ============================================================
	# MODEL LOADING + TOKEN ID SAFETY CHECKS
	# ============================================================

	def load_model(model_name):
	global _current_model_id, _current_model, _current_tokenizer

	if model_name not in MODELS:
	raise ValueError(f"Unknown model: {model_name}")

	model_id = MODELS[model_name]
	if _current_model_id == model_id and _current_model is not None:
	return _current_tokenizer, _current_model

	# Release the previous model before loading another one.
	if _current_model is not None:
	del _current_model
	_current_model = None
	_current_tokenizer = None
	_current_model_id = None
	gc.collect()
	if torch.cuda.is_available():
	torch.cuda.empty_cache()

	print(f"Loading tokenizer: {model_id}")
	tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
	config = AutoConfig.from_pretrained(model_id)

	# Important: if the tokenizer has more IDs than the model config,
	# set the config vocabulary size BEFORE loading weights. This avoids
	# the common 32000-vs-32001 embedding shape mismatch when the checkpoint
	# actually contains the extra embedding row.
	tokenizer_size = len(tokenizer)
	if tokenizer_size > int(config.vocab_size):
	print(f"Adjusting config vocab_size: {config.vocab_size} -> {tokenizer_size}")
	config.vocab_size = tokenizer_size

	model = AutoModelForCausalLM.from_pretrained(
	model_id,
	config=config,
	torch_dtype=DTYPE,
	low_cpu_mem_usage=True,
	)

	embedding_count = model.get_input_embeddings().num_embeddings
	if tokenizer_size > embedding_count:
	# Defensive fallback for architectures whose embedding table was not
	# sized from config. New rows are initialized; see note in UI errors.
	model.resize_token_embeddings(tokenizer_size)
	embedding_count = model.get_input_embeddings().num_embeddings

	# Validate all tokenizer IDs before the first CUDA operation.
	vocab = tokenizer.get_vocab()
	invalid = [(token, idx) for token, idx in vocab.items() if idx < 0 or idx >= embedding_count]
	if invalid:
	del model
	gc.collect()
	raise ValueError(
	f"Tokenizer/model vocabulary mismatch for {model_name}: {len(invalid)} token IDs "
	f"are outside embedding range [0, {embedding_count - 1}]. "
	"Check that this model repository's tokenizer and weights come from the same revision."
	)

	if tokenizer.pad_token_id is None:
	if tokenizer.eos_token_id is not None:
	tokenizer.pad_token = tokenizer.eos_token
	else:
	tokenizer.add_special_tokens({"pad_token": "<\|pad\|>"})
	if len(tokenizer) > model.get_input_embeddings().num_embeddings:
	model.resize_token_embeddings(len(tokenizer))

	model.to(DEVICE)
	model.eval()

	_current_tokenizer = tokenizer
	_current_model = model
	_current_model_id = model_id
	print(
	f"Model ready: {model_id} on {DEVICE} \| tokenizer={len(tokenizer)} \| "
	f"embeddings={model.get_input_embeddings().num_embeddings}"
	)
	return tokenizer, model


	# ============================================================
	# OUTPUT FORMATTING
	# ============================================================

	CODE_BLOCK_PATTERN = re.compile(r"```([^\n`])\n(.?)```", re.DOTALL)


	def format_output(text):
	def replace_code_block(match):
	language = match.group(1).strip() or "text"
	code = match.group(2).rstrip("\n")
	longest_run = max((len(run) for run in re.findall(r"`+", code)), default=0)
	fence = "`" * max(3, longest_run + 1)
	return f"\n\n### `<{language.upper()}>`\n\n{fence}{language}\n{code}\n{fence}\n\n"

	return CODE_BLOCK_PATTERN.sub(replace_code_block, text).strip() or "The model returned an empty response."


	# Preserve the original marker cleanup behavior, but only when those markers
	# actually appear in the model output.
	def normalize_token_markers(text):
	if "Ġ" in text or "Ċ" in text:
	text = text.replace(" ", "")
	text = text.replace("Ġ", " ").replace("Ċ", "\n")
	return text


	# ============================================================
	# GENERATION
	# ============================================================

	@spaces.GPU
	def generate(prompt, model_name, top_p, top_k, temperature, max_new_tokens):
	if not prompt or not prompt.strip():
	yield "Please enter a message."
	return
	if model_name not in MODELS:
	yield "Please select a valid model."
	return

	limits = get_limits(model_name)
	# Clamp all UI values again server-side; browser inputs are not trusted.
	top_p = max(0.1, min(float(top_p), limits["top_p_max"]))
	top_k = max(1, min(int(top_k), limits["top_k_max"]))
	temperature = max(0.1, min(float(temperature), limits["temperature_max"]))
	max_new_tokens = max(1, min(int(max_new_tokens), limits["max_new_tokens_max"]))

	yield "Loading model and generating..."

	try:
	# Keep the lock for the complete operation; only one request/model
	# generation runs at once, avoiding unsafe concurrent CUDA use.
	with _model_lock:
	tokenizer, model = load_model(model_name)
	messages = [{"role": "user", "content": prompt.strip()}]

	if getattr(tokenizer, "chat_template", None):
	inputs = tokenizer.apply_chat_template(
	messages,
	tokenize=True,
	add_generation_prompt=True,
	return_tensors="pt",
	return_dict=True,
	)
	else:
	inputs = tokenizer(prompt.strip(), return_tensors="pt")

	inputs = {key: value.to(DEVICE) for key, value in inputs.items()}
	input_ids = inputs.get("input_ids")
	if input_ids is None or input_ids.numel() == 0:
	raise ValueError("Tokenizer returned empty input_ids.")

	embedding_count = model.get_input_embeddings().num_embeddings
	min_id = int(input_ids.min().item())
	max_id = int(input_ids.max().item())
	if min_id < 0 or max_id >= embedding_count:
	raise ValueError(
	f"Input token ID range [{min_id}, {max_id}] exceeds model embedding range "
	f"[0, {embedding_count - 1}]. Tokenizer and model files are incompatible."
	)

	model_context = getattr(model.config, "max_position_embeddings", None)
	if model_context is None:
	model_context = getattr(model.config, "n_positions", None)
	if isinstance(model_context, int) and model_context > 1:
	# Keep the prompt inside the model context and leave room for generation.
	# This also prevents position-embedding index errors on very long prompts.
	max_input_tokens = model_context - 1
	if input_ids.shape[-1] > max_input_tokens:
	for key in ("input_ids", "attention_mask", "token_type_ids"):
	if key in inputs and inputs[key].ndim == 2:
	inputs[key] = inputs[key][:, -max_input_tokens:]
	input_ids = inputs["input_ids"]
	print(
	f"Prompt truncated to {max_input_tokens} tokens for model context "
	f"limit {model_context}."
	)
	max_new_tokens = min(max_new_tokens, max(1, model_context - input_ids.shape[-1]))

	pad_token_id = tokenizer.pad_token_id
	if pad_token_id is None:
	pad_token_id = tokenizer.eos_token_id

	generation_kwargs = {
	"max_new_tokens": max_new_tokens,
	"do_sample": True,
	"temperature": temperature,
	"top_p": top_p,
	"top_k": top_k,
	"repetition_penalty": REPETITION_PENALTY,
	"pad_token_id": pad_token_id,
	"use_cache": True,
	}
	if tokenizer.eos_token_id is not None:
	generation_kwargs["eos_token_id"] = tokenizer.eos_token_id

	streamer = TextIteratorStreamer(
	tokenizer,
	skip_prompt=True,
	skip_special_tokens=True,
	timeout=0.5,
	)
	worker_errors = []

	def generation_worker():
	try:
	with torch.inference_mode():
	model.generate(inputs, streamer=streamer, generation_kwargs)
	except Exception as exc:
	worker_errors.append(exc)
	try:
	streamer.on_finalized_text("", stream_end=True)
	except Exception:
	pass

	worker = threading.Thread(target=generation_worker, daemon=True)
	worker.start()
	raw_response = ""

	while True:
	try:
	chunk = next(streamer)
	except StopIteration:
	break
	except queue.Empty:
	if not worker.is_alive():
	if worker_errors:
	raise worker_errors[0]
	break
	continue

	if chunk:
	raw_response += chunk
	yield format_output(normalize_token_markers(raw_response))

	worker.join()
	if worker_errors:
	raise worker_errors[0]
	yield format_output(normalize_token_markers(raw_response))

	except torch.cuda.OutOfMemoryError:
	gc.collect()
	if torch.cuda.is_available():
	torch.cuda.empty_cache()
	yield "GPU memory is insufficient. Select a smaller model or reduce Max New Tokens."
	except Exception as exc:
	print("Generation error:", repr(exc))
	traceback.print_exc()
	message = str(exc)
	if "device-side assert" in message.lower() or "cuda error" in message.lower():
	yield (
	"CUDA generation failed. The app checked token IDs before generation, so this may be "
	"a model/checkpoint incompatibility or a CUDA kernel issue. Restart the Space runtime "
	"and verify that the model weights and tokenizer are from the same repository revision."
	)
	else:
	yield f"Generation error ({type(exc).__name__}): {message}"


	# ============================================================
	# DARK THEME
	# ============================================================

	CSS = """
	.gradio-container { max-width: 940px !important; margin: 0 auto !important; background: #0b0d12 !important; }
	#app-header { text-align: center; padding: 22px 0 8px 0; }
	#app-header h1 { font-size: 30px; font-weight: 750; letter-spacing: -1px; color: #f1f5f9; }
	#app-subtitle { text-align: center; color: #929bad; margin-bottom: 22px; }
	#model-selector, #settings-panel { border: 1px solid #292e3b; border-radius: 12px; background: #11141c; }
	#input-area textarea { background: #11141c !important; border: 1px solid #292e3b !important; border-radius: 12px !important; color: #f1f5f9 !important; font-size: 15px !important; line-height: 1.7 !important; }
	#output-area { border: 1px solid #292e3b; border-radius: 12px; background: #11141c; padding: 18px; }
	#output-area h3 { color: #a5b4fc; font-size: 12px; font-weight: 700; letter-spacing: 1px; margin-bottom: 10px; }
	#output-area pre { border: 1px solid #292e3b !important; border-radius: 10px !important; background: #090b10 !important; padding: 16px !important; overflow-x: auto !important; }
	#output-area code { font-size: 13px !important; line-height: 1.7 !important; }
	#generate-button { border-radius: 10px; min-height: 46px; font-weight: 650; }
	footer { display: none !important; }
	"""

	with gr.Blocks(
	theme=gr.themes.Base(
	primary_hue="indigo",
	secondary_hue="slate",
	neutral_hue="slate",
	).set(
	body_background_fill="#0b0d12",
	body_text_color="#e5e7eb",
	block_background_fill="#11141c",
	block_border_color="#292e3b",
	input_background_fill="#11141c",
	button_primary_background_fill="#6366f1",
	button_primary_background_fill_hover="#4f46e5",
	),
	css=CSS,
	title="LLM Playground",
	) as demo:
	gr.Markdown("# LLM Playground", elem_id="app-header")
	gr.Markdown("A simple, modern interface for language models.", elem_id="app-subtitle")

	model_selector = gr.Dropdown(
	choices=list(MODELS.keys()),
	value=DEFAULT_MODEL,
	label="MODEL",
	info="Choose which model to run.",
	elem_id="model-selector",
	interactive=True,
	)

	prompt_input = gr.Textbox(
	label="INPUT",
	placeholder="Ask a question, write code, or enter a prompt...",
	lines=5,
	max_lines=12,
	elem_id="input-area",
	)

	with gr.Accordion("Generation settings", open=True, elem_id="settings-panel"):
	with gr.Row():
	top_p_input = gr.Slider(
	minimum=0.1, maximum=get_limits(DEFAULT_MODEL)["top_p_max"], value=DEFAULTS["top_p"],
	step=0.05, label="Top P", info="Minimum 0.1; maximum depends on selected model."
	)
	top_k_input = gr.Slider(
	minimum=1, maximum=get_limits(DEFAULT_MODEL)["top_k_max"], value=DEFAULTS["top_k"],
	step=1, label="Top K", info="Minimum 1; maximum depends on selected model."
	)
	with gr.Row():
	temperature_input = gr.Slider(
	minimum=0.1, maximum=get_limits(DEFAULT_MODEL)["temperature_max"], value=DEFAULTS["temperature"],
	step=0.05, label="Temperature", info="Minimum 0.1; maximum depends on selected model."
	)
	max_new_tokens_input = gr.Slider(
	minimum=1, maximum=get_limits(DEFAULT_MODEL)["max_new_tokens_max"], value=DEFAULTS["max_new_tokens"],
	step=1, label="Max New Tokens", info="Minimum 1; maximum depends on selected model."
	)

	generate_button = gr.Button("Generate", variant="primary", elem_id="generate-button")
	output = gr.Markdown(value="Your model's response will appear here.", label="OUTPUT", elem_id="output-area")

	model_selector.change(
	fn=update_model_controls,
	inputs=[model_selector],
	outputs=[top_p_input, top_k_input, temperature_input, max_new_tokens_input],
	)

	generation_inputs = [
	prompt_input, model_selector, top_p_input, top_k_input,
	temperature_input, max_new_tokens_input,
	]
	generate_button.click(fn=generate, inputs=generation_inputs, outputs=output)
	prompt_input.submit(fn=generate, inputs=generation_inputs, outputs=output)

	if __name__ == "__main__":
	demo.queue(default_concurrency_limit=1).launch(
	server_name="0.0.0.0",
	server_port=7860,
	)