Spaces:
Sleeping
Sleeping
File size: 5,341 Bytes
4e316d6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | """
HuggingFace Transformers generation engine — drop-in replacement for
GenerationEngine that loads any model from the HuggingFace Hub.
Uses ``AutoModelForCausalLM`` and ``AutoTokenizer`` so any causal LM
works without architecture-specific code. Exposes the same
``generate_answer()`` / ``generate_stream()`` interface so eval scripts
can swap engines with a single line.
Designed for the multi-model eval benchmark: load a model, run the
eval, delete the engine, free memory, load the next model.
"""
from __future__ import annotations
import re
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from src.utils.device import get_default_device
# Pattern to strip <think>...</think> blocks (possibly spanning lines)
_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)
# Unclosed <think> block at the end of a generation
_THINK_UNCLOSED_RE = re.compile(r"<think>.*", re.DOTALL)
class TransformersEngine:
"""Architecture-agnostic generation engine backed by HuggingFace Transformers.
Parameters
----------
model_id : str
HuggingFace model identifier (e.g. ``"Qwen/Qwen3-4B"``).
dtype : torch.dtype
Weight precision. Defaults to ``torch.bfloat16``.
max_memory_gb : float | None
Optional memory cap (not enforced, just for documentation).
"""
def __init__(
self,
model_id: str,
dtype: torch.dtype = torch.bfloat16,
enable_thinking: bool = False,
trust_remote_code: bool = True,
):
self.model_id = model_id
self.device = get_default_device()
self.enable_thinking = enable_thinking
self.tokenizer = AutoTokenizer.from_pretrained(
model_id, trust_remote_code=trust_remote_code,
)
self.model = AutoModelForCausalLM.from_pretrained(
model_id,
dtype=dtype,
trust_remote_code=trust_remote_code,
).to(self.device)
self.model.eval()
def generate_answer(
self,
messages: list[dict],
max_tokens: int = 512,
temperature: float = 0.7,
top_k: int = 20,
top_p: float = 0.95,
) -> str:
"""Generate a complete response (non-streaming).
Compatible with ``GenerationEngine.generate_answer()``.
"""
# Apply chat template
template_kwargs = {
"tokenize": False,
"add_generation_prompt": True,
}
# Disable thinking mode when supported (Qwen3/3.5) to avoid
# wasting the token budget on <think> blocks.
if not self.enable_thinking:
try:
text = self.tokenizer.apply_chat_template(
messages, enable_thinking=False, **template_kwargs,
)
except TypeError:
# Model doesn't support enable_thinking param
text = self.tokenizer.apply_chat_template(
messages, **template_kwargs,
)
else:
text = self.tokenizer.apply_chat_template(
messages, **template_kwargs,
)
inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
input_len = inputs["input_ids"].shape[1]
# Generate — fall back to greedy if sampling hits NaN logits
# (common on MPS + bfloat16 with long contexts).
do_sample = temperature > 0.01
gen_kwargs: dict = dict(
max_new_tokens=max_tokens,
pad_token_id=self.tokenizer.eos_token_id,
)
if do_sample:
gen_kwargs.update(
do_sample=True,
temperature=max(temperature, 1e-4),
top_k=top_k,
top_p=top_p,
)
else:
gen_kwargs["do_sample"] = False
with torch.no_grad():
try:
outputs = self.model.generate(**inputs, **gen_kwargs)
except RuntimeError as e:
if "probability tensor" in str(e) and do_sample:
# NaN/Inf logits — retry with greedy decoding
outputs = self.model.generate(
**inputs,
max_new_tokens=max_tokens,
pad_token_id=self.tokenizer.eos_token_id,
do_sample=False,
)
else:
raise
# Decode only the generated tokens
generated = outputs[0][input_len:]
answer = self.tokenizer.decode(generated, skip_special_tokens=True)
# Strip <think>...</think> blocks
answer = _THINK_RE.sub("", answer)
answer = _THINK_UNCLOSED_RE.sub("", answer)
return answer.strip()
def generate_stream(
self,
messages: list[dict],
max_tokens: int = 512,
temperature: float = 0.7,
top_k: int = 20,
top_p: float = 0.95,
):
"""Yield the full answer as a single chunk.
Provided for interface compatibility with ``GenerationEngine``.
True token-level streaming is not needed for eval.
"""
yield self.generate_answer(
messages,
max_tokens=max_tokens,
temperature=temperature,
top_k=top_k,
top_p=top_p,
)
|