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=[])