app.py

import os
os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8')
os.environ.setdefault('TOKENIZERS_PARALLELISM', 'false')
os.environ.setdefault('GGML_CUDA_DISABLE_GRAPHS', '1')

import spaces
import ctypes
import glob
import json
import site
import time
from pathlib import Path

import gradio as gr
import numpy as np
import torch
import transformers
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer

from highlight import token_cost, text_spans
from ui import build_demo, theme, CSS

from codec import (
    ArithmeticDecoder, ArithmeticEncoder, MAX_BYTES, MAX_TOKENS,
    logits_to_cdf, pack, profile_id, unpack, verify_text,
)

PROFILE = json.loads(Path(__file__).with_name('profile.json').read_text())
tokenizer = AutoTokenizer.from_pretrained(PROFILE['model_id'], revision=PROFILE['revision'])
print('Downloading published Qwen3.5 Base i1-Q4_K_M GGUF, unchanged.', flush=True)
MODEL_PATH = hf_hub_download(PROFILE['gguf_repo'], PROFILE['gguf_file'], revision=PROFILE['gguf_revision'])
print('GGUF ready; GPU initialization happens on the first request.', flush=True)
_LLM = None


def get_model():
    # llama.cpp allocates CUDA directly, so initialize only within @spaces.GPU.
    global _LLM
    started = time.monotonic()
    if _LLM is None:
        for library in ['libcudart.so*', 'libcublasLt.so*', 'libcublas.so*']:
            for base in site.getsitepackages():
                for path in glob.glob(str(Path(base) / 'nvidia' / '*' / 'lib' / library)):
                    ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)
        import llama_cpp
        if not torch.cuda.is_available() or not llama_cpp.llama_supports_gpu_offload():
            raise RuntimeError('The CUDA backend is unavailable; this demo requires GPU execution.')
        _LLM = llama_cpp.Llama(
            model_path=MODEL_PATH, n_gpu_layers=-1,
            n_ctx=PROFILE['n_ctx'], n_batch=1, n_ubatch=1,
            n_threads=2, n_threads_batch=2,
            logits_all=False, flash_attn=False,
            # ggml.h enum; this binding does not export the BF16 constant.
            type_k=30, type_v=30,
            use_mmap=True, verbose=True,
        )
        print(f'Qwen GPU model loaded in {time.monotonic()-started:.2f}s.', flush=True)
    return _LLM, round(time.monotonic()-started, 2)


def text_tokens(text):
    if not isinstance(text, str):
        raise ValueError('Enter text to compress.')
    if len(text.encode('utf-8', errors='strict')) > MAX_BYTES:
        raise ValueError('Keep the input under 64 KiB.')
    ids = tokenizer.encode(text, add_special_tokens=False)
    if len(ids) > MAX_TOKENS:
        raise ValueError(f'This demo accepts 8,192 tokens; your text has {len(ids):,}.')
    if detokenize(ids) != text:
        raise ValueError('This text is changed by the model tokenizer, so it cannot be compressed losslessly here.')
    return ids


def detokenize(ids):
    return tokenizer.decode(ids, skip_special_tokens=False, clean_up_tokenization_spaces=False) if ids else ''


def runtime_profile():
    import llama_cpp
    settings = {
        **PROFILE, 'llama_cpp': llama_cpp.__version__, 'torch': torch.__version__,
        'transformers': transformers.__version__, 'numpy': np.__version__,
        'cuda': torch.version.cuda, 'gpu': torch.cuda.get_device_name(0),
        'capability': list(torch.cuda.get_device_capability(0)),
        'cuda_graphs': False, 'cpu_threads': 2,
    }
    return profile_id(settings), settings


def next_distribution(llm, token):
    llm.eval([token])
    # Read every logit; never sample or truncate to top-k probabilities.
    logits = np.ctypeslib.as_array(llm._ctx.get_logits_ith(-1), shape=(llm.n_vocab(),))
    return logits_to_cdf(logits)


def encode_tokens(llm, ids, progress):
    llm.reset()  # clears both attention cache and the recurrent DeltaNet state
    encoder = ArithmeticEncoder()
    costs = []
    for index, token in enumerate(ids):
        cdf = next_distribution(llm, ids[index-1] if index else PROFILE['prefix_token_id'])
        encoder.write(token, cdf)
        costs.append(token_cost(cdf, token))
        if index % 4 == 0 or index == len(ids)-1:
            progress((index+1, len(ids)), desc='Encoding', unit='tokens')
    return encoder.finish(), costs


def decode_packet(llm, packet, profile, progress):
    if packet.profile != profile:
        raise ValueError('This code requires its original model/runtime/GPU profile. The current worker differs.')
    llm.reset()
    decoder, ids = ArithmeticDecoder(packet.payload), []
    for index in range(packet.token_count):
        cdf = next_distribution(llm, ids[-1] if ids else PROFILE['prefix_token_id'])
        ids.append(decoder.read(cdf))
        if index % 4 == 0 or index == packet.token_count-1:
            progress((index+1, packet.token_count), desc='Restoring', unit='tokens')
    text = detokenize(ids)
    verify_text(text, packet)
    return text


# Measured on ZeroGPU: ~3.5 s load and ~10 ms per token; the first task after a restart
# adds ~24 s of one-time CUDA warm-up. spaces scales these by 1.5 on this GPU type.
def compress_duration(text, progress=None):
    return int(min(360, 25 + 0.012*len(text_tokens(text))))


def decompress_duration(code, progress=None):
    return int(min(360, 25 + 0.012*unpack(code).token_count))


@spaces.GPU(duration=compress_duration, size='large')
def compress(text, progress=gr.Progress()):
    try:
        ids = text_tokens(text)
        progress(None, desc='Loading model onto GPU…')
        llm, load_seconds = get_model()
        if llm.tokenize(text.encode('utf-8'), add_bos=False, special=True) != ids:
            raise ValueError('The GGUF and source tokenizers disagree on this input.')
        profile, runtime = runtime_profile()
        start = time.monotonic()
        payload, costs = encode_tokens(llm, ids, progress)
        encoded_at = time.monotonic()
        code = pack(payload, text, len(ids), profile)
        packet = unpack(code)
        original = packet.original_bytes
        stats = {
            'verified': False, 'original_utf8_bytes': original, 'tokens': len(ids),
            'token_spans': text_spans(text, ids, costs, [llm.detokenize([t], special=True) for t in ids]),
            'arithmetic_bytes': len(payload), 'packet_bytes': packet.packet_bytes,
            'ascii_characters_and_bytes': len(code),
            'size_percent': round(100*len(code)/original, 1) if original else None,
            'encode_seconds': round(encoded_at-start, 2),
            'verification_seconds': 0,
            'gpu_load_seconds': load_seconds, 'quantization': PROFILE['quantization'],
            'gpu': runtime['gpu'], 'profile': profile.hex(),
            'timing_note': 'Encode excludes GPU loading and queue. Automatic decode verification is off.',
        }
        ratio = f'{100*len(code)/original:.1f}% of original size' if original else 'empty input'
        report = (f'**{original:,} → {len(code):,} bytes** · {ratio} · {len(ids):,} tokens · {stats["encode_seconds"]}s encode')
        return code, report, stats
    except (ValueError, RuntimeError) as exc:
        raise gr.Error(str(exc)) from exc


@spaces.GPU(duration=decompress_duration, size='large')
def decompress(code, progress=gr.Progress()):
    try:
        packet = unpack(code)
        progress(None, desc='Loading model onto GPU…')
        llm, load_seconds = get_model()
        profile, runtime = runtime_profile()
        start = time.monotonic()
        text = decode_packet(llm, packet, profile, progress)
        elapsed = time.monotonic()-start
        return text, f'**Checksum verified.** Restored {packet.original_bytes:,} bytes in {elapsed:.2f}s.', {
            'verified': True, 'decode_seconds': round(elapsed, 2), 'gpu_load_seconds': load_seconds,
            'original_utf8_bytes': packet.original_bytes, 'tokens': packet.token_count,
            'gpu': runtime['gpu'], 'profile': profile.hex(),
        }
    except (ValueError, RuntimeError) as exc:
        raise gr.Error(str(exc)) from exc


demo = build_demo(compress, decompress)

if __name__ == '__main__':
    demo.queue(max_size=8).launch(css=CSS, theme=theme(), show_error=True, footer_links=[])