| """BFB232 prototype inference. MIT license; see LICENSE and THIRD_PARTY_NOTICES.md.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import bisect |
| import contextlib |
| import json |
| import math |
| import os |
| from pathlib import Path |
| import threading |
| import time |
|
|
| |
| os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") |
|
|
| import torch |
| from safetensors.torch import load_file |
| from transformers import AutoModel, AutoTokenizer |
|
|
| import bfb232_text as textio |
|
|
|
|
| class BoundaryHead(torch.nn.Module): |
| """Frozen B187/B232 'left' binary head; CLOSE score = logit[1] - logit[0].""" |
|
|
| def __init__(self): |
| super().__init__() |
| self.base = torch.nn.Sequential( |
| torch.nn.LayerNorm(768), torch.nn.Linear(768, 96), |
| torch.nn.GELU(), torch.nn.Linear(96, 2), |
| ) |
| self.extra = torch.nn.Sequential( |
| torch.nn.LayerNorm(384), torch.nn.Linear(384, 64), |
| torch.nn.GELU(), torch.nn.Linear(64, 2), |
| ) |
|
|
| def forward(self, features): |
| return self.base(features[:, :768]) + self.extra(features[:, 768:1152]) |
|
|
|
|
| def make_window(tokenizer, text, positions, left, right): |
| """Exact bfb130 indexed token anchors and 256-character side pools.""" |
| encoding = tokenizer(text[left:right], return_offsets_mapping=True, |
| add_special_tokens=True, truncation=False) |
| offsets = encoding.pop("offset_mapping") |
| if len(encoding["input_ids"]) > 8192: |
| raise ValueError("Retokenized window exceeds the reference 8192-token limit") |
| valid = [i for i, (a, b) in enumerate(offsets) if b > a] |
| starts = [offsets[i][0] for i in valid] |
| if starts != sorted(starts): |
| raise ValueError("Non-monotone tokenizer offsets") |
| anchors, after, sides = [], [], [] |
| for position in positions: |
| pos = position - left |
| at = bisect.bisect_left(starts, pos) |
| if at == 0: |
| raise ValueError("Candidate has no preceding token") |
| a = valid[at - 1] |
| b = valid[at] if at < len(valid) else a |
| lo = bisect.bisect_left(starts, max(0, pos - 256)) |
| hi = bisect.bisect_left(starts, pos + 256) |
| anchors.append(a) |
| after.append(b) |
| sides.append((valid[lo:at] or [a], valid[at:hi] or [b])) |
| return dict(encoding=dict(encoding), anchors=anchors, after=after, |
| sides=sides, left=left, right=right) |
|
|
|
|
| def inference_windows(tokenizer, text, positions): |
| """Canonical text/positions; 7900-token lattice, stride 6500, interior owner.""" |
| encoded = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True, |
| truncation=False, verbose=False) |
| offsets = encoded["offset_mapping"] |
| if not positions: |
| return [], len(offsets) |
| if len(offsets) <= 7900: |
| w = make_window(tokenizer, text, positions, 0, len(text)) |
| w["indices"] = list(range(len(positions))) |
| return [w], len(offsets) |
| starts = sorted(set(list(range(0, len(offsets) - 7900 + 1, 6500)) |
| + [len(offsets) - 7900])) |
| spans = [(offsets[s][0] if s else 0, |
| offsets[s + 7899][1] if s + 7900 < len(offsets) else len(text)) |
| for s in starts] |
| groups = {} |
| for i, p in enumerate(positions): |
| choices = [j for j, (a, b) in enumerate(spans) if a < p <= b] |
| if not choices: |
| raise ValueError("Candidate lies outside the token windows") |
| owner = max(choices, key=lambda j: min(p - spans[j][0], spans[j][1] - p)) |
| groups.setdefault(owner, []).append(i) |
| result = [] |
| for j, indices in sorted(groups.items()): |
| w = make_window(tokenizer, text, [positions[i] for i in indices], *spans[j]) |
| w["indices"] = indices |
| result.append(w) |
| return result, len(offsets) |
|
|
|
|
| def features_from_hidden(hidden, window): |
| """Preserves reference FP32 pooling and operation order.""" |
| pair = torch.cat([hidden[window["anchors"]], hidden[window["after"]]], -1).float() |
| left = torch.stack([hidden[a].float().mean(0) for a, _ in window["sides"]]) |
| right = torch.stack([hidden[b].float().mean(0) for _, b in window["sides"]]) |
| return torch.cat([pair, left, right, right - left, right * left], -1) |
|
|
|
|
| class BFB232: |
| def __init__(self, tokenizer, encoder, head, config, device): |
| self.tokenizer, self.encoder, self.head = tokenizer, encoder, head |
| self.config, self.device = config, torch.device(device) |
| self.thresholds = dict(config["thresholds"]) |
| self.lock = threading.Lock() |
|
|
| @classmethod |
| def from_pretrained(cls, model_path, *, device="auto", revision=None, |
| local_files_only=False, token=None): |
| """Load a local release folder or Hub ID; never executes Hub Python code. |
| |
| For a Hub ID, pin revision to its commit hash. CPU uses FP32; the |
| reference CUDA path uses BF16 autocast with FP32 stored weights. |
| """ |
| folder = Path(model_path) |
| if not folder.is_dir(): |
| from huggingface_hub import snapshot_download |
| folder = Path(snapshot_download( |
| repo_id=str(model_path), revision=revision, token=token, |
| local_files_only=local_files_only, |
| allow_patterns=["bfb232_config.json", "encoder/*.json", |
| "encoder/*.safetensors", "tokenizer/*.json", |
| "boundary_head.safetensors"], |
| )) |
| config = json.loads((folder / "bfb232_config.json").read_text("utf-8")) |
| if config["format_version"] != 1 or config["input_contract"] != textio.VERSION: |
| raise ValueError("Unsupported BFB232 release/input contract") |
| if device == "auto": |
| device = "cuda" if torch.cuda.is_available() and torch.cuda.is_bf16_supported(including_emulation=False) else "cpu" |
| device = torch.device(device) |
| if device.type not in ("cpu", "cuda"): |
| raise ValueError("Only CPU FP32 and CUDA BF16 are supported") |
| if device.type == "cuda": |
| if not torch.cuda.is_available(): |
| raise ValueError("CUDA is unavailable; use device='cpu'") |
| with torch.cuda.device(device): |
| if not torch.cuda.is_bf16_supported(including_emulation=False): |
| raise ValueError("CUDA BF16 support is required; use device='cpu'") |
| torch.set_num_threads(4) |
| torch.use_deterministic_algorithms(True) |
| torch.backends.cudnn.benchmark = False |
| torch.backends.cudnn.deterministic = True |
| torch.backends.cuda.matmul.allow_tf32 = False |
| torch.backends.cudnn.allow_tf32 = False |
| tokenizer = AutoTokenizer.from_pretrained(folder / "tokenizer", local_files_only=True) |
| encoder, info = AutoModel.from_pretrained( |
| folder / "encoder", local_files_only=True, use_safetensors=True, |
| dtype=torch.float32, attn_implementation="sdpa", output_loading_info=True, |
| ) |
| if any(info.get(k) for k in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")): |
| raise ValueError(f"Encoder state mismatch: {info}") |
| head = BoundaryHead() |
| head.load_state_dict(load_file(str(folder / "boundary_head.safetensors")), strict=True) |
| encoder.to(device).eval().requires_grad_(False) |
| head.to(device).eval().requires_grad_(False) |
| return cls(tokenizer, encoder, head, config, device) |
|
|
| def segment(self, text, language): |
| """Return source-preserving segments and original Unicode code-point offsets. |
| |
| Language must be 'ko' or 'en'; no automatic language detection or |
| minimum/maximum segment length is imposed. Empty input returns []. |
| """ |
| if not isinstance(text, str): |
| raise TypeError("text must be a string") |
| if language not in self.thresholds: |
| raise ValueError("language must be 'ko' or 'en'") |
| with self.lock: |
| return self._segment(text, language) |
|
|
| def _segment(self, text, language): |
| started = time.perf_counter() |
| canonical = textio.normalize(text) |
| positions = textio.candidates(canonical.text, language) |
| original_positions = [canonical.to_original[p] for p in positions] |
| windows, input_tokens = inference_windows(self.tokenizer, canonical.text, positions) |
| scores = [None] * len(positions) |
| processed_tokens = 0 |
| autocast = (torch.autocast("cuda", dtype=torch.bfloat16) |
| if self.device.type == "cuda" else contextlib.nullcontext()) |
| with torch.inference_mode(), autocast: |
| for w in windows: |
| inputs = self.tokenizer.pad([w["encoding"]], padding=True, return_tensors="pt") |
| inputs = {k: v.to(self.device) for k, v in inputs.items()} |
| hidden = self.encoder(**inputs).last_hidden_state[0] |
| logits = self.head(features_from_hidden(hidden, w)) |
| values = (logits[:, 1] - logits[:, 0]).float().cpu().tolist() |
| for i, score in zip(w["indices"], values): |
| if scores[i] is not None or not math.isfinite(score): |
| raise RuntimeError("Invalid/duplicate boundary score") |
| scores[i] = score |
| processed_tokens += len(w["encoding"]["input_ids"]) |
| if any(s is None for s in scores): |
| raise RuntimeError("Missing boundary scores") |
| cuts = [p for p, s in zip(original_positions, scores) if s >= self.thresholds[language]] |
| segments, left = [], 0 |
| for right in cuts + ([len(text)] if text else []): |
| segments.append(dict(index=len(segments) + 1, start=left, end=right, |
| text=text[left:right], boundary="END" if right == len(text) else "CLOSE")) |
| left = right |
| if "".join(s["text"] for s in segments) != text: |
| raise RuntimeError("Source preservation failed") |
| return dict(model="BFB232", language=language, segments=segments, |
| source_preserved=True, input_contract=textio.VERSION, |
| candidate_positions=original_positions, scores=scores, |
| threshold=self.thresholds[language], |
| precision="cuda-bf16" if self.device.type == "cuda" else "cpu-fp32", |
| metrics=dict(input_chars=len(text), input_tokens=input_tokens, |
| processed_tokens=processed_tokens, windows=len(windows), |
| candidates=len(positions), segment_count=len(segments), |
| total_seconds=time.perf_counter() - started)) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="BFB232 semantic segmentation") |
| parser.add_argument("--model", default=str(Path(__file__).parent)) |
| parser.add_argument("--revision", help="Hub commit hash when --model is a Hub ID") |
| parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"]) |
| parser.add_argument("--language", required=True, choices=["ko", "en"]) |
| parser.add_argument("--input", required=True, type=Path, help="UTF-8 text file") |
| parser.add_argument("--output", required=True, type=Path, help="New JSON file; refuses overwrite") |
| args = parser.parse_args() |
| if args.output.exists(): |
| parser.error("Output exists; choose a new path") |
| |
| with args.input.open(encoding="utf-8", newline="") as source: |
| text = source.read() |
| model = BFB232.from_pretrained(args.model, device=args.device, revision=args.revision) |
| result = model.segment(text, args.language) |
| with args.output.open("x", encoding="utf-8", newline="\n") as out: |
| json.dump(result, out, ensure_ascii=False, indent=2) |
| out.write("\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
