"""Model registry and generation helpers. A 9B model in bf16 needs ~18 GB, so several of them will not fit on a single A10G (24 GB) at once. We therefore keep a *single-slot* cache: whenever a different model is requested we free the previous one before loading the new one. Single-correction traffic (always AMALIA) reuses the cached weights; the benchmark cycles through models one at a time. Hardware is detected at load time: - CUDA with >= 20 GB VRAM (A10G/ZeroGPU): full precision (bf16, or fp16 on pre-Ampere cards that lack bf16 support). - CUDA with < 20 GB VRAM (e.g. an RTX 2060 6 GB): 4-bit nf4 quantization, with any overflow layers offloaded to CPU RAM. - No CUDA: plain CPU load - only viable for small models with plenty of RAM. Override with LOAD_IN_4BIT=1 (force 4-bit) or LOAD_IN_4BIT=0 (force full precision) regardless of the detected VRAM. """ import gc import os from threading import Thread import torch from transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer, ) # Display name -> Hugging Face repo id. MODELS = { "AMALIA-9B": "amalia-llm/AMALIA-9B-0626-SFT", "EuroLLM-9B": "utter-project/EuroLLM-9B-Instruct", "Llama-3.1-8B": "meta-llama/Llama-3.1-8B-Instruct", "Gemma-2-9B": "google/gemma-2-9b-it", "Qwen2.5-3B": "Qwen/Qwen2.5-3B-Instruct", # small enough for 6 GB cards "Tucano-2B": "TucanoBR/Tucano-2b4-Instruct", } # The specialised pt-PT model we want to showcase. CHAMPION = "AMALIA-9B" # Same base/size as the champion but without pt-PT tuning: the fair baseline. # The IPT gap between CHAMPION and REFERENCE isolates the effect of the tuning. REFERENCE = "EuroLLM-9B" # None lets huggingface_hub fall back to the token stored by `hf auth login`. _HF_TOKEN = os.getenv("HF_TOKEN") # Single-slot cache: {"name", "tokenizer", "model"}. _current = {"name": None, "tokenizer": None, "model": None} # Timeout (s) for the streamer queue: if generation stalls this long the UI # gets an exception instead of hanging forever. _STREAM_TIMEOUT = 600.0 # --------------------------------------------------------------------------- # Hardware detection # --------------------------------------------------------------------------- def _supports_bf16(): return torch.cuda.get_device_capability(0)[0] >= 8 # Ampere or newer def hardware_summary(): """One-line description of what inference will run on (for the UI).""" # On ZeroGPU the CUDA device only exists inside @spaces.GPU functions, # so don't touch torch.cuda at UI-build time. if os.getenv("SPACE_ID"): return "Hugging Face Space (ZeroGPU, bf16)" if not torch.cuda.is_available(): return "CPU (sem CUDA) — apenas modelos pequenos são viáveis" props = torch.cuda.get_device_properties(0) total = props.total_memory / 1024**3 mode = "4-bit nf4" if _use_4bit() else ("bf16" if _supports_bf16() else "fp16") return f"{props.name} ({total:.0f} GB VRAM) — {mode}" def _use_4bit(): forced = os.getenv("LOAD_IN_4BIT") if forced == "1": return True if forced == "0": return False total = torch.cuda.get_device_properties(0).total_memory / 1024**3 return total < 20 # a 9B in 16-bit needs ~18 GB just for weights def _max_memory(): """Cap the GPU budget at the VRAM that is *actually free* right now, minus headroom for the KV cache and CUDA overhead. No CPU budget is given on purpose: spilling 4-bit models to CPU RAM keeps the spilled layers in fp32 (8x larger), which on a nearly-full machine kills the process with a native access violation in torch_cpu.dll. With quantized CPU offload disallowed, a model that does not fit fails at the *planning* stage with a catchable ValueError before any weight is read. """ vram_free, _ = torch.cuda.mem_get_info() gpu_budget = max(vram_free - 0.6 * 1024**3, 512 * 1024**2) return {0: int(gpu_budget)} def _load_kwargs(): """Build from_pretrained kwargs for the detected hardware.""" if not torch.cuda.is_available(): # float32 on CPU: most CPUs have no fast bf16 path and fp16 CPU # generation is numerically fragile. return dict(dtype=torch.float32, low_cpu_mem_usage=True) compute_dtype = torch.bfloat16 if _supports_bf16() else torch.float16 if not _use_4bit(): return dict(dtype=compute_dtype, device_map="auto") quant = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=compute_dtype, ) return dict( dtype=compute_dtype, device_map="auto", max_memory=_max_memory(), quantization_config=quant, ) def _free(): """Drop the currently loaded model and reclaim GPU/CPU memory.""" _current["model"] = None _current["tokenizer"] = None _current["name"] = None gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def get_model(name): """Return (tokenizer, model) for ``name``, loading it if necessary.""" if _current["name"] == name and _current["model"] is not None: return _current["tokenizer"], _current["model"] _free() repo_id = MODELS[name] tokenizer = AutoTokenizer.from_pretrained(repo_id, token=_HF_TOKEN) try: model = AutoModelForCausalLM.from_pretrained( repo_id, token=_HF_TOKEN, **_load_kwargs() ) except ValueError as exc: if "offload" not in str(exc).lower(): raise # The device-map plan needed CPU/disk spill: the model does not fit # in the VRAM that is free right now. Fail before reading any weight. vram_free, _ = torch.cuda.mem_get_info() raise RuntimeError( f"'{name}' não cabe na VRAM livre neste momento " f"({vram_free / 1024**3:.1f} GB). Fecha aplicações que usem a " "placa gráfica, escolhe um modelo mais pequeno (ex.: Qwen2.5-3B) " "ou usa a versão no Hugging Face Space." ) from exc # If accelerate ran out of RAM/VRAM it silently leaves weights on the # "meta" device and generation later fails with a cryptic error # ("Tensor.item() cannot be called on meta tensors"). Fail loudly instead. if any(p.is_meta for p in model.parameters()): _free() raise RuntimeError( f"'{name}' não coube na memória disponível (pesos ficaram em " "'meta'). Numa máquina local, define LOAD_IN_4BIT=1 e garante que " "o torch tem CUDA; ou escolhe um modelo mais pequeno (Qwen2.5-3B)." ) model.eval() if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token _current.update(name=name, tokenizer=tokenizer, model=model) return tokenizer, model def _apply_template(tokenizer, system, user_text): """Build input ids, folding the system prompt into the user turn for models (e.g. Gemma-2) whose chat template rejects a dedicated system role.""" messages = [ {"role": "system", "content": system}, {"role": "user", "content": user_text}, ] try: ids = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt" ) except Exception: merged = f"{system}\n\n{user_text}" messages = [{"role": "user", "content": merged}] ids = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt" ) # In newer transformers, apply_chat_template may return a BatchEncoding # (dict-like) instead of a plain tensor. Extract the tensor if needed. if not isinstance(ids, torch.Tensor): ids = ids["input_ids"] return ids def _budget_new_tokens(tokenizer, user_text, cap): """A corrector's output should have roughly the input's length; capping at ~2x the input stops runaway degeneration on weak models.""" n_input = len(tokenizer(user_text, add_special_tokens=False)["input_ids"]) return max(96, min(cap, 2 * n_input + 128)) def _generation_kwargs(tokenizer, model, input_ids, max_new_tokens): return dict( input_ids=input_ids, attention_mask=torch.ones_like(input_ids), max_new_tokens=max_new_tokens, do_sample=False, # deterministic: a proofreader must not "invent" repetition_penalty=1.05, pad_token_id=tokenizer.pad_token_id, # Some models (e.g. Tucano) imitate their template tags instead of # emitting EOS; the "<<<" guards against echoing our own delimiters. stop_strings=["", "", "<<<"], tokenizer=tokenizer, ) def start_stream(name, system, user_text, max_new_tokens=1024): """Generate in a background thread; return an iterator of text chunks. Exceptions raised inside the generation thread (OOM, gated repo, ...) are re-raised in the consumer instead of hanging the stream forever. """ tokenizer, model = get_model(name) input_ids = _apply_template(tokenizer, system, user_text).to(model.device) max_new_tokens = _budget_new_tokens(tokenizer, user_text, max_new_tokens) streamer = TextIteratorStreamer( tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=_STREAM_TIMEOUT, ) errors = [] def _worker(): try: model.generate( streamer=streamer, **_generation_kwargs(tokenizer, model, input_ids, max_new_tokens), ) except Exception as exc: # surfaced to the consumer below errors.append(exc) streamer.end() # unblock the iterator Thread(target=_worker, daemon=True).start() def _iterate(): for chunk in streamer: yield chunk if errors: raise errors[0] return _iterate() @torch.no_grad() def generate_text(name, system, user_text, max_new_tokens=512): """Blocking generation that returns the full decoded completion.""" tokenizer, model = get_model(name) input_ids = _apply_template(tokenizer, system, user_text).to(model.device) max_new_tokens = _budget_new_tokens(tokenizer, user_text, max_new_tokens) output = model.generate( **_generation_kwargs(tokenizer, model, input_ids, max_new_tokens) ) completion = output[0, input_ids.shape[-1]:] return tokenizer.decode(completion, skip_special_tokens=True).strip()