File size: 8,331 Bytes
317d342 | 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 | """Dedicated-process, explicitly opted-in development scoring; no text generation."""
import importlib.metadata
import json
import math
import os
from pathlib import Path
from . import _windows
from .artifact import SOURCE_KIND_MODE, WINDOW_GATE_MODE, inspect_artifact
from .protocol import encode_response, validate_request
ENVIRONMENT = {
'torch': '2.11.0+cu128', 'transformers': '5.19.0', 'safetensors': '0.8.0',
'tokenizers': '0.23.2', 'huggingface_hub': '1.33.0', 'numpy': '2.5.3',
}
def initialize_windows_system_root():
if os.name == 'nt':
import ctypes
buffer = ctypes.create_unicode_buffer(32768)
get = ctypes.WinDLL('kernel32', use_last_error=True).GetWindowsDirectoryW
get.argtypes = [ctypes.c_wchar_p, ctypes.c_uint]
get.restype = ctypes.c_uint
size = get(buffer, len(buffer))
if not 0 < size < len(buffer):
raise ValueError('windows_directory_unavailable')
os.environ['SystemRoot'] = buffer.value
class DevelopmentScorer:
"""Opt-in prototype. Use in its own process; GPU parity is still required."""
def __init__(self, artifact, revision, *, cache_directory, development=False, precision='fp32',
whole_source_context=False):
if development is not True or precision not in ('fp32', 'fp16'):
raise ValueError('explicit_development_precision_required')
if type(whole_source_context) is not bool or (whole_source_context and precision != 'fp32'):
raise ValueError('explicit_FP32_whole_source_context_required')
self.whole_source_context = whole_source_context
self.root = Path(artifact)
self.report = inspect_artifact(self.root, revision)
if self.report['mode'] in (WINDOW_GATE_MODE, SOURCE_KIND_MODE) and precision != 'fp32':
raise ValueError('window_gate_FP32_required')
self.root = self.root.resolve()
self.revision, self.precision = revision, precision
self.cache = Path(cache_directory)
if not self.cache.is_absolute() or self.cache.resolve().is_relative_to(self.root):
raise ValueError('separate_absolute_cache_directory_required')
self._tokenizer = self._model = self._torch = None
self._token_config = None
self.last_metrics = None
def _prepare_environment(self):
initialization = self.report.get('initialization_recipe')
declared = initialization.get('environment', []) if isinstance(initialization, dict) else []
if (not isinstance(declared, list) or len(declared) != len(ENVIRONMENT)
or set(declared) != {name + '==' + value for name, value in ENVIRONMENT.items()}):
raise ValueError('supported_provisioned_environment_required')
for name, expected in ENVIRONMENT.items():
if importlib.metadata.version(name) != expected:
raise ValueError('supported_provisioned_environment_required')
initialize_windows_system_root()
self.cache.mkdir(parents=True, exist_ok=True)
os.environ.update(HF_HUB_OFFLINE='1', TRANSFORMERS_OFFLINE='1',
HF_HOME=str(self.cache / 'hf'), TORCHINDUCTOR_CACHE_DIR=str(self.cache / 'torch'),
TOKENIZERS_PARALLELISM='false', CUBLAS_WORKSPACE_CONFIG=':4096:8')
def _load_tokenizer(self):
self._prepare_environment()
from tokenizers import Tokenizer
self._tokenizer = Tokenizer.from_file(str(self.root / 'tokenizer/tokenizer.json'))
self._token_config = json.loads((self.root / 'tokenizer/tokenizer_config.json').read_bytes())
def _load_model(self):
self.report = inspect_artifact(self.root, self.revision)
if self.report['mode'] in (WINDOW_GATE_MODE, SOURCE_KIND_MODE) and self.precision != 'fp32':
raise ValueError('window_gate_FP32_required')
import torch
if (not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 0)
or torch.cuda.get_device_name() != 'NVIDIA GeForce RTX 5080'):
raise ValueError('tested_local_CUDA_target_required_no_fallback')
from transformers.initialization import no_init_weights
from transformers.models.modernbert.configuration_modernbert import ModernBertConfig
from transformers.models.modernbert.modeling_modernbert import ModernBertModel
from ._model import JointGroupModel
model_type = JointGroupModel
if self.report['mode'] == WINDOW_GATE_MODE:
from ._window_gate import WindowGateModel
model_type = WindowGateModel
elif self.report['mode'] == SOURCE_KIND_MODE:
from ._source_kind import SourceKindGateModel
model_type = SourceKindGateModel
torch.set_num_threads(4)
torch.set_num_interop_threads(1)
torch.use_deterministic_algorithms(True)
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
torch.backends.cudnn.benchmark = False
config = ModernBertConfig.from_pretrained(self.root, local_files_only=True)
config._attn_implementation = 'sdpa'
with no_init_weights():
model = model_type(ModernBertModel(config))
state = torch.load(self.root / 'checkpoint.pt', map_location='cpu', weights_only=True, mmap=True)
if any(not isinstance(value, torch.Tensor) or value.dtype != torch.float32
or not torch.isfinite(value).all() for value in state.values()):
raise ValueError('finite_FP32_source_tensors_required')
model.load_state_dict(state, strict=True, assign=True)
dtype = torch.float32 if self.precision == 'fp32' else torch.float16
self._model = model.to(device='cuda', dtype=dtype).eval()
self._torch = torch
def score(self, request):
"""Admit every original group before loading; return no partial response."""
self.last_metrics = None
request = validate_request(request, self.revision)
if self._tokenizer is None:
self._load_tokenizer()
windows, packing = _windows.pack_windows(self._tokenizer, self._token_config, request)
if not self.whole_source_context and (len(windows) != 1 or packing['context_tokens_per_window']):
raise ValueError('single_complete_window_required_no_truncation')
owned = [i for _, ranges, indices in windows for i in indices]
if (not 1 <= len(windows) <= _windows.MAX_WINDOWS
or any(not indices or len(ranges) != len(indices) for _, ranges, indices in windows)
or any(type(i) is not int for i in owned)
or sorted(owned) != list(range(len(request['groups'])))):
raise ValueError('score_alignment')
if not self.whole_source_context and owned != list(range(len(request['groups']))):
raise ValueError('score_alignment')
if self._model is None:
self._load_model()
torch = self._torch
dtype = torch.float32 if self.precision == 'fp32' else torch.float16
parameters = list(self._model.parameters())
if not parameters or any(p.device.type != 'cuda' or p.dtype != dtype for p in parameters):
raise ValueError('declared_CUDA_precision_required_no_fallback')
scores = [None] * len(request['groups'])
for ids, ranges, indices in windows:
inputs = torch.tensor([ids], device='cuda')
with torch.inference_mode():
values = self._model(inputs, torch.ones_like(inputs), ranges).cpu().tolist()
if (not isinstance(values, list) or len(values) != len(indices)
or any(type(v) not in (int, float) or not math.isfinite(v) for v in values)):
raise ValueError('invalid_scores')
for i, value in zip(indices, values):
scores[i] = value
encoded = encode_response(request, scores)
self.last_metrics = dict(quality='unqualified', qualification=False, encoder_forwards=len(windows),
device='cuda', precision=self.precision, packing=packing)
if self.whole_source_context:
self.last_metrics['whole_source_context'] = True
return json.loads(encoded)
|