KernelKind's picture
Link the published research code and dataset
eb32c81 verified
|
Raw History Blame Contribute Delete
6.26 kB

Extract subjects from one post

Choose a model-family folder below. Download one checkpoint directory and model.json from that folder, plus the shared ASv6.100.md from the repository root. Read that family's checkpoints.json to identify the checkpoint and verify its weight/configuration hashes. The numbers in checkpoint names count training-note uses; they are not optimizer steps.

The original runtime pins Python 3.12, PyTorch 2.12.1+cu130, Transformers 5.17.0, PEFT 0.21.0, Accelerate 1.15.0 and safetensors 0.8.0. Qwen CUDA execution also uses fla-core 0.5.2. Use a CUDA installation compatible with your hardware; these pins describe the research runtime rather than a promise of compatibility with all GPU hosts. The code below follows its text-backbone loading and decoding path. A fresh GPU execution of this standalone example has not been tested.

The published adapters have FP32 weights. Evaluation loaded an unquantized BF16 base, applied the adapter and merged it into the in-memory evaluation copy. The original experiments used 96 GB GPUs. Smaller-memory and quantized inference have not been validated. The adapter download size excludes the base weights, host memory during loading and generation cache.

Put the files in a local directory with this layout, and save one post as UTF-8 in post.txt:

subject-model/
  ASv6.100.md
  post.txt
  qwen3.8-27b/            # or gemma-4-31b-it / gemma-4-26b-a4b-it
    model.json
    checkpoint-1000/      # or another checkpoint from that family's checkpoints.json
      adapter_config.json
      adapter_model.safetensors

Save the following as extract_subjects.py inside subject-model/. Set family and checkpoint to the downloaded model-family/checkpoint folders. The example reads local adapter files, downloads only the pinned upstream base/tokenizer if needed, and makes no model API calls.

from pathlib import Path
import hashlib
import json
import torch
from peft import PeftModel
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from torch.nn.attention import SDPBackend, sdpa_kernel

root = Path(__file__).resolve().parent
family = root / "qwen3.8-27b"
settings = json.loads((family / "model.json").read_text())
checkpoint = family / "checkpoint-1000"
base, revision = settings["base_model"], settings["base_revision"]
instruction = (root / "ASv6.100.md").read_text(encoding="utf-8").rstrip()
if hashlib.sha256(instruction.encode()).hexdigest() != settings["instruction_text_sha256"]:
    raise ValueError("ASv6.100 instruction does not match training")
post = (root / "post.txt").read_text(encoding="utf-8")
messages = [{"role": "user", "content": instruction + "\n\n<post>\n" + post + "\n</post>"}]

tokenizer = AutoTokenizer.from_pretrained(base, revision=revision, trust_remote_code=False)
if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token
architecture = AutoConfig.from_pretrained(base, revision=revision, trust_remote_code=False)
options = dict(revision=revision, trust_remote_code=False, dtype=torch.bfloat16,
               attn_implementation="sdpa", config=architecture.get_text_config())
if architecture.model_type in {"gemma4", "gemma4_unified"}:
    options["key_mapping"] = {r"^model\.language_model\.": "model."}
if architecture.model_type in {"qwen3_5", "qwen3_5_text"}:
    from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
model, loading = AutoModelForCausalLM.from_pretrained(base, output_loading_info=True, **options)
if loading["missing_keys"] or loading["mismatched_keys"] or loading.get("error_msgs"):
    raise ValueError("Pinned base weights did not load completely")
model = model.to("cuda")
model = PeftModel.from_pretrained(model, checkpoint, is_trainable=False)
model = model.merge_and_unload(safe_merge=True).eval()

prompt = tokenizer.apply_chat_template(messages, tokenize=False,
                                      add_generation_prompt=True, enable_thinking=False)
tokens = tokenizer.encode(prompt, add_special_tokens=False)
allowance = settings["generation"]["max_new_tokens"]
if len(tokens) + allowance > settings["generation"]["max_length"]:
    raise ValueError("Post and output allowance exceed the recorded context limit")
inputs = torch.tensor([tokens], device="cuda")
eos = model.generation_config.eos_token_id
if eos is None:
    eos = tokenizer.eos_token_id
stop_ids = [eos] if isinstance(eos, int) else eos
with torch.inference_mode(), sdpa_kernel([SDPBackend.FLASH_ATTENTION,
                                        SDPBackend.EFFICIENT_ATTENTION]):
    output = model.generate(input_ids=inputs, attention_mask=torch.ones_like(inputs),
                            do_sample=False, max_new_tokens=allowance, use_cache=True,
                            pad_token_id=tokenizer.pad_token_id, eos_token_id=eos)
answer = output[0, len(tokens):].tolist()
stop = next((i for i, token in enumerate(answer) if token in stop_ids), None)
answer = answer if stop is None else answer[:stop + 1]
raw = tokenizer.decode(answer, skip_special_tokens=True)
print(raw, flush=True)
if stop is None:
    raise ValueError("Output reached its limit without an end-of-sequence token")

From that directory:

python extract_subjects.py

The recorded decoding uses thinking off, greedy generation, a 20,480-token context and up to 10,240 new tokens. Do not silently truncate a post or the instruction to fit. Output is JSON containing a subjects list, with number, context and subject fields. Dotted numbers represent children within this post's local subject hierarchy. Returning valid JSON does not establish that the subjects are correct. The research reproduction guide uses the project loader and parser and covers training.

Only subject extraction is covered here. Definition writing, vector encoding and construction of a hierarchy across many posts are separate stages.