fragment-2 v3.0.2 — from-scratch release (auto-published by the training notebook)

40bf0a1 verified 29 days ago

12 kB

#!/usr/bin/env python3
# f3.py — self-contained Python runtime for FrameXlabs/fragment-2 (v3).
#
# Loads the released safetensors checkpoint + BPE tokenizer and answers typed
# questions with calibrated probabilities in a single forward pass.
# Five question types: choice / score / noul / multi / rank.
#
# Usage:
#     from f3 import Fragment
#     m = Fragment.from_pretrained("FrameXlabs/fragment-2")  # hub or local dir
#     res = m.decide(state="...", questions={...})
#
# Dependencies: torch, safetensors.
import json
import math
import os
import re
from typing import Dict

import torch
import torch.nn as nn
import torch.nn.functional as F


class F3Tokenizer:
    # Portable char-level BPE (the fragment build, proven since fragment-1).
    # Non-Latin scripts fall back to single characters.

    def __init__(self, path):
        spec = json.load(open(path))
        self.vocab = spec["vocab"]
        self.id_of = {v: i for i, v in enumerate(self.vocab)}
        self.ranks = {(a, b): i for i, (a, b) in enumerate(spec["merges"])}
        self.added = set(spec.get("added", []))
        self.pad = self.id_of["[PAD]"]
        self.cls = self.id_of["[CLS]"]
        self.sep = self.id_of["[SEP]"]
        self.unk = self.id_of["[UNK]"]
        self._cache = {}

    def _bpe(self, word):
        syms = list(word)
        while len(syms) >= 2:
            best, rank = None, None
            for a, b in zip(syms, syms[1:]):
                r = self.ranks.get((a, b))
                if r is not None and (rank is None or r < rank):
                    best, rank = (a, b), r
            if best is None:
                break
            a, b = best
            out, j = [], 0
            while j < len(syms):
                if j < len(syms) - 1 and syms[j] == a and syms[j + 1] == b:
                    out.append(a + b)
                    j += 2
                else:
                    out.append(syms[j])
                    j += 1
            syms = out
        return syms

    def _word_ids(self, word):
        cached = self._cache.get(word)
        if cached is not None:
            return cached
        ids = [self.id_of.get(s, self.unk) for s in self._bpe(word)]
        if len(self._cache) < 300_000:
            self._cache[word] = ids
        return ids

    def encode(self, text):
        text = text.lower()
        ids = []
        for w in re.findall(r"@?[a-z0-9]+|[^\sa-z0-9]", text):
            if w in self.added:
                ids.append(self.id_of[w])
            else:
                ids.extend(self._word_ids(w))
        return ids

class RMSNorm(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.w = nn.Parameter(torch.ones(d))

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * self.w


class RoPE(nn.Module):
    # Rotary position embeddings (theta=10000). Parameter-free.
    def __init__(self, d, max_len, base=10000.0):
        super().__init__()
        inv = 1.0 / (base ** (torch.arange(0, d, 2).float() / d))
        t = torch.arange(max_len).float()
        freqs = torch.outer(t, inv)                     # (L, d/2)
        self.register_buffer("cos", freqs.cos(), persistent=False)
        self.register_buffer("sin", freqs.sin(), persistent=False)

    def forward(self, x):                              # x: (B, H, L, D)
        L = x.shape[-2]
        cos, sin = self.cos[:L], self.sin[:L]           # (L, D/2)
        x1, x2 = x[..., ::2], x[..., 1::2]              # even/odd split
        cos = cos[None, None, :, :]
        sin = sin[None, None, :, :]
        out = torch.empty_like(x)
        out[..., ::2] = x1 * cos - x2 * sin
        out[..., 1::2] = x1 * sin + x2 * cos
        return out


class Block(nn.Module):
    # Pre-norm transformer block: MHA (RoPE) + SwiGLU FFN. No bias.
    def __init__(self, d, n_heads, ff):
        super().__init__()
        self.n1 = RMSNorm(d)
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.o = nn.Linear(d, d, bias=False)
        self.n2 = RMSNorm(d)
        self.gate = nn.Linear(d, ff, bias=False)
        self.up = nn.Linear(d, ff, bias=False)
        self.down = nn.Linear(ff, d, bias=False)
        self.d, self.h = d, n_heads
        self.rope = None

    def forward(self, x, att):
        # att: (B, L) True = real token (used as a broadcastable key mask)
        B, L, _ = x.shape
        h = self.n1(x)
        q, k, v = self.qkv(h).chunk(3, -1)
        q = q.view(B, L, self.h, -1).transpose(1, 2)    # (B,H,L,dh)
        k = k.view(B, L, self.h, -1).transpose(1, 2)
        v = v.view(B, L, self.h, -1).transpose(1, 2)
        q, k = self.rope(q), self.rope(k)
        a = F.scaled_dot_product_attention(q, k, v,
                                           attn_mask=att[:, None, None, :])
        x = x + self.o(a.transpose(1, 2).reshape(B, L, -1))
        h = self.n2(x)
        return x + self.down(F.silu(self.gate(h)) * self.up(h))


class F3Net(nn.Module):
    # Fragment-2 encoder: the original F2Net, scaled and modernized.
    # 10 layers, d=384, 6 heads, SwiGLU ff=1536, RoPE, RMSNorm, max_len 256.
    # Per-position scalar logit head (the proven fragment-1 design).
    def __init__(self, vocab=32768, d=384, max_len=256, n_layers=10,
                 n_heads=6, ff=1536):
        super().__init__()
        self.d, self.max_len = d, max_len
        self.scale = math.sqrt(d)
        self.tok = nn.Embedding(vocab, d, padding_idx=0)
        self.blocks = nn.ModuleList([Block(d, n_heads, ff) for _ in range(n_layers)])
        self.ln_f = RMSNorm(d)
        self.h1 = nn.Linear(d, d)
        self.h2 = nn.Linear(d, 1)
        self.rope = None
        self.cfg = {"vocab": vocab, "d": d, "max_len": max_len,
                    "n_layers": n_layers, "n_heads": n_heads, "ff": ff}
        self.apply(self._init)
        nn.init.normal_(self.tok.weight, std=0.02)
        nn.init.constant_(self.tok.weight[0], 0.0)

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    def _init_rope(self, device):
        # RoPE is parameter-free; created lazily and shared across blocks.
        # NOTE: rotations act on head_dim (d / n_heads), not model dim.
        if self.rope is None or self.rope.cos.device != device:
            r = RoPE(self.d // self.cfg["n_heads"], self.max_len).to(device)
            for blk in self.blocks:
                blk.rope = r
            self.rope = r

    def forward(self, ids):
        self._init_rope(ids.device)
        att = ids != 0
        B, L = ids.shape
        h = self.tok(ids) * self.scale
        for blk in self.blocks:
            h = blk(h, att)
        h = self.ln_f(h)
        return self.h2(F.gelu(self.h1(h))).squeeze(-1).float()


MAX_LEN = 256


def _confidence(p):
    k = len(p)
    if k < 2:
        return 1.0
    ent = -sum(x * math.log(max(x, 1e-12)) for x in p)
    return min(1.0, max(0.0, 1.0 - ent / math.log(k)))


class Fragment:
    def __init__(self, root):
        self.cfg = json.load(open(os.path.join(root, "f3_config.json")))
        self.tok = F3Tokenizer(os.path.join(root, "tokenizer.json"))
        from safetensors.torch import load_file
        sd = load_file(os.path.join(root, "model.safetensors"))
        sd = {k: (v.float() if v.dtype is not torch.float32 else v)
              for k, v in sd.items()}
        arch = self.cfg["architecture"]
        self.net = F3Net(arch["vocab"], arch["d"], arch["max_len"],
                         arch["n_layers"], arch["n_heads"], arch["ff"])
        self.net.load_state_dict(sd, strict=True)
        self.net.eval()
        self.version = self.cfg.get("version", "3.0")
        self.temps = self.cfg.get("temperatures", {})

    @classmethod
    def from_pretrained(cls, src="FrameXlabs/fragment-2"):
        if os.path.isdir(src):
            return cls(src)
        from huggingface_hub import snapshot_download
        return cls(snapshot_download(repo_id=src))

    def _options(self, q):
        t = q["type"]
        if t == "choice":
            crit = q.get("criteria") or {}
            return [f"@{i} {k}: {v}" if v else f"@{i} {k}"
                    for i, (k, v) in enumerate(crit.items())], list(crit.keys())
        if t in ("score", "multi", "rank"):
            crit = q.get("criteria") or []
            return [f"@{i} {c}" for i, c in enumerate(crit)], None
        return ["@no: no, the statement does not hold",
                "@yes: yes, the statement holds"], None

    @torch.no_grad()
    def decide(self, state, questions):
        import time as _t
        t0 = _t.time()
        answers = {}
        for qid, q in questions.items():
            opts, keys = self._options(q)
            pre = self.tok.encode(f"{q['instructions']} input:")
            opt_ids = self.tok.encode("options: " + "; ".join(opts))
            markers = [o.split()[0].rstrip(":") for o in opts]
            marker_ids = [self.tok.id_of[m] for m in markers]
            positions, seen = [], set()
            for idx, t in enumerate(opt_ids):
                if t in marker_ids and t not in seen:
                    seen.add(t)
                    positions.append(idx)
            budget = MAX_LEN - 2 - len(pre) - len(opt_ids)
            body = self.tok.encode(" " + state)[: max(8, budget)]
            ids = [self.tok.cls] + pre + body + opt_ids + [self.tok.sep]
            ids = ids[:MAX_LEN]
            shift = 1 + len(pre) + len(body)
            abs_pos = [p + shift for p in positions]
            abs_pos = [p for p in abs_pos if p < len(ids) - 1]
            t = torch.tensor([ids])
            logits = self.net(t)[0, abs_pos]
            temp = self.temps.get(q["type"], 1.0) or 1.0
            names = [(keys[i] if keys else f"option_{i}")
                     for i in range(len(positions))]
            base = {"type": q["type"], "question": q["instructions"]}
            if q["type"] == "choice":
                p = torch.softmax(logits / temp, -1).tolist()
                base["probabilities"] = {names[i]: round(p[i], 4)
                                         for i in range(len(p))}
                base["choice"] = names[int(max(range(len(p)),
                                               key=lambda i: p[i]))]
                base["confidence"] = round(_confidence(p), 4)
            elif q["type"] == "score":
                p = torch.softmax(logits / temp, -1).tolist()
                exp = sum(i * p[i] for i in range(len(p)))
                base["probabilities"] = {str(i): round(p[i], 4)
                                         for i in range(len(p))}
                base["score"] = round(exp, 4)
            elif q["type"] == "multi":
                p = torch.sigmoid(logits / temp).tolist()
                base["probabilities"] = {names[i]: round(p[i], 4)
                                         for i in range(len(p))}
                base["selected"] = [names[i] for i in range(len(p))
                                    if p[i] > 0.5]
            elif q["type"] == "rank":
                z = (logits / temp).tolist()
                order = sorted(range(len(z)), key=lambda i: -z[i])
                p = torch.softmax(logits / temp, -1).tolist()
                base["probabilities"] = {names[i]: round(p[i], 4)
                                         for i in range(len(p))}
                base["order"] = [names[i] for i in order]
            else:  # noul
                p = torch.softmax(logits / temp, -1).tolist()
                base["probabilities"] = {"false": round(p[0], 4),
                                         "true": round(p[1], 4)}
                base["noul"] = round(p[1], 4)
            answers[qid] = base
        return {"model": "fragment-2", "version": self.version,
                "answers": answers,
                "usage": {"latency_ms": round((_t.time() - t0) * 1000)}}

	#!/usr/bin/env python3
	# f3.py — self-contained Python runtime for FrameXlabs/fragment-2 (v3).
	#
	# Loads the released safetensors checkpoint + BPE tokenizer and answers typed
	# questions with calibrated probabilities in a single forward pass.
	# Five question types: choice / score / noul / multi / rank.
	#
	# Usage:
	# from f3 import Fragment
	# m = Fragment.from_pretrained("FrameXlabs/fragment-2") # hub or local dir
	# res = m.decide(state="...", questions={...})
	#
	# Dependencies: torch, safetensors.
	import json
	import math
	import os
	import re
	from typing import Dict

	import torch
	import torch.nn as nn
	import torch.nn.functional as F


	class F3Tokenizer:
	# Portable char-level BPE (the fragment build, proven since fragment-1).
	# Non-Latin scripts fall back to single characters.

	def __init__(self, path):
	spec = json.load(open(path))
	self.vocab = spec["vocab"]
	self.id_of = {v: i for i, v in enumerate(self.vocab)}
	self.ranks = {(a, b): i for i, (a, b) in enumerate(spec["merges"])}
	self.added = set(spec.get("added", []))
	self.pad = self.id_of["[PAD]"]
	self.cls = self.id_of["[CLS]"]
	self.sep = self.id_of["[SEP]"]
	self.unk = self.id_of["[UNK]"]
	self._cache = {}

	def _bpe(self, word):
	syms = list(word)
	while len(syms) >= 2:
	best, rank = None, None
	for a, b in zip(syms, syms[1:]):
	r = self.ranks.get((a, b))
	if r is not None and (rank is None or r < rank):
	best, rank = (a, b), r
	if best is None:
	break
	a, b = best
	out, j = [], 0
	while j < len(syms):
	if j < len(syms) - 1 and syms[j] == a and syms[j + 1] == b:
	out.append(a + b)
	j += 2
	else:
	out.append(syms[j])
	j += 1
	syms = out
	return syms

	def _word_ids(self, word):
	cached = self._cache.get(word)
	if cached is not None:
	return cached
	ids = [self.id_of.get(s, self.unk) for s in self._bpe(word)]
	if len(self._cache) < 300_000:
	self._cache[word] = ids
	return ids

	def encode(self, text):
	text = text.lower()
	ids = []
	for w in re.findall(r"@?[a-z0-9]+\|[^\sa-z0-9]", text):
	if w in self.added:
	ids.append(self.id_of[w])
	else:
	ids.extend(self._word_ids(w))
	return ids

	class RMSNorm(nn.Module):
	def __init__(self, d):
	super().__init__()
	self.w = nn.Parameter(torch.ones(d))

	def forward(self, x):
	return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * self.w


	class RoPE(nn.Module):
	# Rotary position embeddings (theta=10000). Parameter-free.
	def __init__(self, d, max_len, base=10000.0):
	super().__init__()
	inv = 1.0 / (base ** (torch.arange(0, d, 2).float() / d))
	t = torch.arange(max_len).float()
	freqs = torch.outer(t, inv) # (L, d/2)
	self.register_buffer("cos", freqs.cos(), persistent=False)
	self.register_buffer("sin", freqs.sin(), persistent=False)

	def forward(self, x): # x: (B, H, L, D)
	L = x.shape[-2]
	cos, sin = self.cos[:L], self.sin[:L] # (L, D/2)
	x1, x2 = x[..., ::2], x[..., 1::2] # even/odd split
	cos = cos[None, None, :, :]
	sin = sin[None, None, :, :]
	out = torch.empty_like(x)
	out[..., ::2] = x1 * cos - x2 * sin
	out[..., 1::2] = x1 * sin + x2 * cos
	return out


	class Block(nn.Module):
	# Pre-norm transformer block: MHA (RoPE) + SwiGLU FFN. No bias.
	def __init__(self, d, n_heads, ff):
	super().__init__()
	self.n1 = RMSNorm(d)
	self.qkv = nn.Linear(d, 3 * d, bias=False)
	self.o = nn.Linear(d, d, bias=False)
	self.n2 = RMSNorm(d)
	self.gate = nn.Linear(d, ff, bias=False)
	self.up = nn.Linear(d, ff, bias=False)
	self.down = nn.Linear(ff, d, bias=False)
	self.d, self.h = d, n_heads
	self.rope = None

	def forward(self, x, att):
	# att: (B, L) True = real token (used as a broadcastable key mask)
	B, L, _ = x.shape
	h = self.n1(x)
	q, k, v = self.qkv(h).chunk(3, -1)
	q = q.view(B, L, self.h, -1).transpose(1, 2) # (B,H,L,dh)
	k = k.view(B, L, self.h, -1).transpose(1, 2)
	v = v.view(B, L, self.h, -1).transpose(1, 2)
	q, k = self.rope(q), self.rope(k)
	a = F.scaled_dot_product_attention(q, k, v,
	attn_mask=att[:, None, None, :])
	x = x + self.o(a.transpose(1, 2).reshape(B, L, -1))
	h = self.n2(x)
	return x + self.down(F.silu(self.gate(h)) * self.up(h))


	class F3Net(nn.Module):
	# Fragment-2 encoder: the original F2Net, scaled and modernized.
	# 10 layers, d=384, 6 heads, SwiGLU ff=1536, RoPE, RMSNorm, max_len 256.
	# Per-position scalar logit head (the proven fragment-1 design).
	def __init__(self, vocab=32768, d=384, max_len=256, n_layers=10,
	n_heads=6, ff=1536):
	super().__init__()
	self.d, self.max_len = d, max_len
	self.scale = math.sqrt(d)
	self.tok = nn.Embedding(vocab, d, padding_idx=0)
	self.blocks = nn.ModuleList([Block(d, n_heads, ff) for _ in range(n_layers)])
	self.ln_f = RMSNorm(d)
	self.h1 = nn.Linear(d, d)
	self.h2 = nn.Linear(d, 1)
	self.rope = None
	self.cfg = {"vocab": vocab, "d": d, "max_len": max_len,
	"n_layers": n_layers, "n_heads": n_heads, "ff": ff}
	self.apply(self._init)
	nn.init.normal_(self.tok.weight, std=0.02)
	nn.init.constant_(self.tok.weight[0], 0.0)

	def _init(self, m):
	if isinstance(m, nn.Linear):
	nn.init.normal_(m.weight, std=0.02)
	if m.bias is not None:
	nn.init.zeros_(m.bias)
	elif isinstance(m, nn.Embedding):
	nn.init.normal_(m.weight, std=0.02)

	def _init_rope(self, device):
	# RoPE is parameter-free; created lazily and shared across blocks.
	# NOTE: rotations act on head_dim (d / n_heads), not model dim.
	if self.rope is None or self.rope.cos.device != device:
	r = RoPE(self.d // self.cfg["n_heads"], self.max_len).to(device)
	for blk in self.blocks:
	blk.rope = r
	self.rope = r

	def forward(self, ids):
	self._init_rope(ids.device)
	att = ids != 0
	B, L = ids.shape
	h = self.tok(ids) * self.scale
	for blk in self.blocks:
	h = blk(h, att)
	h = self.ln_f(h)
	return self.h2(F.gelu(self.h1(h))).squeeze(-1).float()


	MAX_LEN = 256


	def _confidence(p):
	k = len(p)
	if k < 2:
	return 1.0
	ent = -sum(x * math.log(max(x, 1e-12)) for x in p)
	return min(1.0, max(0.0, 1.0 - ent / math.log(k)))


	class Fragment:
	def __init__(self, root):
	self.cfg = json.load(open(os.path.join(root, "f3_config.json")))
	self.tok = F3Tokenizer(os.path.join(root, "tokenizer.json"))
	from safetensors.torch import load_file
	sd = load_file(os.path.join(root, "model.safetensors"))
	sd = {k: (v.float() if v.dtype is not torch.float32 else v)
	for k, v in sd.items()}
	arch = self.cfg["architecture"]
	self.net = F3Net(arch["vocab"], arch["d"], arch["max_len"],
	arch["n_layers"], arch["n_heads"], arch["ff"])
	self.net.load_state_dict(sd, strict=True)
	self.net.eval()
	self.version = self.cfg.get("version", "3.0")
	self.temps = self.cfg.get("temperatures", {})

	@classmethod
	def from_pretrained(cls, src="FrameXlabs/fragment-2"):
	if os.path.isdir(src):
	return cls(src)
	from huggingface_hub import snapshot_download
	return cls(snapshot_download(repo_id=src))

	def _options(self, q):
	t = q["type"]
	if t == "choice":
	crit = q.get("criteria") or {}
	return [f"@{i} {k}: {v}" if v else f"@{i} {k}"
	for i, (k, v) in enumerate(crit.items())], list(crit.keys())
	if t in ("score", "multi", "rank"):
	crit = q.get("criteria") or []
	return [f"@{i} {c}" for i, c in enumerate(crit)], None
	return ["@no: no, the statement does not hold",
	"@yes: yes, the statement holds"], None

	@torch.no_grad()
	def decide(self, state, questions):
	import time as _t
	t0 = _t.time()
	answers = {}
	for qid, q in questions.items():
	opts, keys = self._options(q)
	pre = self.tok.encode(f"{q['instructions']} input:")
	opt_ids = self.tok.encode("options: " + "; ".join(opts))
	markers = [o.split()[0].rstrip(":") for o in opts]
	marker_ids = [self.tok.id_of[m] for m in markers]
	positions, seen = [], set()
	for idx, t in enumerate(opt_ids):
	if t in marker_ids and t not in seen:
	seen.add(t)
	positions.append(idx)
	budget = MAX_LEN - 2 - len(pre) - len(opt_ids)
	body = self.tok.encode(" " + state)[: max(8, budget)]
	ids = [self.tok.cls] + pre + body + opt_ids + [self.tok.sep]
	ids = ids[:MAX_LEN]
	shift = 1 + len(pre) + len(body)
	abs_pos = [p + shift for p in positions]
	abs_pos = [p for p in abs_pos if p < len(ids) - 1]
	t = torch.tensor([ids])
	logits = self.net(t)[0, abs_pos]
	temp = self.temps.get(q["type"], 1.0) or 1.0
	names = [(keys[i] if keys else f"option_{i}")
	for i in range(len(positions))]
	base = {"type": q["type"], "question": q["instructions"]}
	if q["type"] == "choice":
	p = torch.softmax(logits / temp, -1).tolist()
	base["probabilities"] = {names[i]: round(p[i], 4)
	for i in range(len(p))}
	base["choice"] = names[int(max(range(len(p)),
	key=lambda i: p[i]))]
	base["confidence"] = round(_confidence(p), 4)
	elif q["type"] == "score":
	p = torch.softmax(logits / temp, -1).tolist()
	exp = sum(i * p[i] for i in range(len(p)))
	base["probabilities"] = {str(i): round(p[i], 4)
	for i in range(len(p))}
	base["score"] = round(exp, 4)
	elif q["type"] == "multi":
	p = torch.sigmoid(logits / temp).tolist()
	base["probabilities"] = {names[i]: round(p[i], 4)
	for i in range(len(p))}
	base["selected"] = [names[i] for i in range(len(p))
	if p[i] > 0.5]
	elif q["type"] == "rank":
	z = (logits / temp).tolist()
	order = sorted(range(len(z)), key=lambda i: -z[i])
	p = torch.softmax(logits / temp, -1).tolist()
	base["probabilities"] = {names[i]: round(p[i], 4)
	for i in range(len(p))}
	base["order"] = [names[i] for i in order]
	else: # noul
	p = torch.softmax(logits / temp, -1).tolist()
	base["probabilities"] = {"false": round(p[0], 4),
	"true": round(p[1], 4)}
	base["noul"] = round(p[1], 4)
	answers[qid] = base
	return {"model": "fragment-2", "version": self.version,
	"answers": answers,
	"usage": {"latency_ms": round((_t.time() - t0) * 1000)}}