| """Scientific-paper classification, public v2 format: LIST FIRST for choices; paper first for yes/no. |
| Apache-2.0. No generated text, no remote model code. The Transformers backend does not reuse prefix states; |
| use vllm_fast.FastPaperClassifier for high-throughput shared-taxonomy batches. |
| """ |
| from __future__ import annotations |
| import json |
| import random |
| from pathlib import Path |
|
|
| SYSTEM = ("You classify scientific papers. You get a paper (its title, and its abstract when available), then a " |
| "question. For a list of options, answer with the code of the option that best describes the research " |
| "content of the paper, following the classification's own meaning of each option. For a single topic, answer " |
| "yes if the topic is a central subject of the paper, else no. On topic: the paper studies the topic itself as " |
| "a main subject: the phenomenon, its causes, consequences, mechanisms, measurement, methods, or how it is " |
| "managed, in any system or organism. Not on topic: the topic is absent or only mentioned in passing; or the " |
| "words mean something else in this paper; or the paper studies a neighbouring subject without studying the " |
| "topic itself.") |
| NONE = "None of these (outside the options listed)" |
| QUESTION = "Which option best describes the research content of this paper?" |
| SCHEME = "Classification of research" |
|
|
|
|
| def paper_text(title, abstract=None, venue=None, year=None): |
| """Benchmark paper representation: title, optional venue/year, abstract limited to 1,500 characters. |
| Title-only inputs are valid. A complete paper is not required or read by this model.""" |
| s = f"Title: {title.strip()}" |
| if venue: |
| s += f"\nVenue: {venue.strip()}" |
| if year: |
| s += f"\nYear: {year}" |
| abstract = " ".join((abstract or "").split()) |
| return s + (f"\nAbstract: {abstract[:1500]}" if abstract else "\n(no abstract)") |
|
|
|
|
| def messages(user): |
| return [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}] |
|
|
|
|
| def option_line(code, o, describe=True, parents=True): |
| extra = [x for x in ((o.get("about") if describe else None), |
| (f"in {o['parent']}" if parents and o.get("parent") else None)) if x] |
| return f"{code}. {o['name']}" + (" — " + "; ".join(extra) if extra else "") |
|
|
|
|
| def head_text(scheme, question, options, codes, describe=True, parents=True): |
| lines = "\n".join(option_line(c, o, describe, parents) for c, o in zip(codes, options)) |
| return f"## Classification: {scheme}\n{question}\nOptions:\n{lines}\n" |
|
|
|
|
| def tail_text(paper): |
| return f"\n## Paper\n{paper}\n\nAnswer with the option code only." |
|
|
|
|
| def yesno_prompt(paper, topic, definition=None): |
| d = f" (what it covers: {definition})" if definition else "" |
| return f"{paper}\n\n## Topic: {topic}{d}\nIs this topic a central subject of the paper? Answer yes or no." |
|
|
|
|
| def read_config(model, revision=None): |
| p = Path(model) / "classifier.json" |
| if not p.exists(): |
| from huggingface_hub import hf_hub_download |
| p = Path(hf_hub_download(model, "classifier.json", revision=revision)) |
| cfg = json.loads(p.read_text()) |
| if cfg.get("format") != "list-first-filler-v1": |
| raise ValueError("This helper is for the v2/list-first release, not the earlier paper-first models.") |
| return cfg |
|
|
|
|
| class PaperInputs: |
| """Exact token construction. No truncation of taxonomy/options or the answer prefix. |
| Long paper text alone is capped at paper_tokens. Oversized prompts raise ValueError instead of dropping labels. |
| The prefix is tokenized once and padded with individual newline token IDs to a cache-block boundary. |
| Do not concatenate these newlines as a string and retokenize: token merging changes the prompt.""" |
| def __init__(self, tokenizer, config, block=None, paper_tokens=700, max_tokens=4000): |
| self.tok, self.cfg = tokenizer, config |
| self.block = int(block or config.get("block", 544)) |
| self.paper_tokens, self.max_tokens = paper_tokens, max_tokens |
| if self.block < 1 or paper_tokens < 1 or max_tokens > 4096: |
| raise ValueError("block/paper_tokens must be positive; this helper supports max_tokens <= 4096") |
| probe = tokenizer.apply_chat_template(messages("XXUSERXX"), tokenize=False, |
| add_generation_prompt=True, enable_thinking=False) |
| self.pre, self.post = probe.split("XXUSERXX") |
| filler = self.encode("\n") |
| if len(filler) != 1 or filler[0] != config["filler"]: |
| raise ValueError("Tokenizer/newline token differs from the release.") |
| self.filler = filler[0] |
| for code, token in zip(config["codes"], config["code_ids"]): |
| if self.encode(code) != [token]: |
| raise ValueError(f"Wrong tokenizer for answer code {code}") |
| for word in ("yes", "no"): |
| if self.encode(word) != [config[word]]: |
| raise ValueError(f"Wrong tokenizer for {word}") |
|
|
| def encode(self, text): |
| return self.tok.encode(text, add_special_tokens=False) |
|
|
| def normalize_options(self, options, none_option=False): |
| out = [] |
| for x in options: |
| o = {"name": x} if isinstance(x, str) else dict(x) |
| if not isinstance(o.get("name"), str) or not o["name"].strip(): |
| raise ValueError("Every option needs a nonempty name") |
| o["id"] = o.get("id", o["name"]) |
| if not isinstance(o["id"], str): |
| raise ValueError("Option IDs must be strings") |
| out.append(o) |
| if none_option: |
| out.append({"name": NONE, "id": "__none__"}) |
| if not 2 <= len(out) <= len(self.cfg["codes"]): |
| raise ValueError(f"Supply 2–{len(self.cfg['codes'])} options; larger taxonomies require multiple questions") |
| if len({o["id"] for o in out}) != len(out): |
| raise ValueError("Option names/IDs must be unique; give each duplicate name a distinct 'id'") |
| return out |
|
|
| def prefix(self, opts, scheme, question, describe=True, parents=True): |
| codes = self.cfg["codes"][:len(opts)] |
| h = self.encode(self.pre + head_text(scheme, question, opts, codes, describe, parents)) |
| return h + [self.filler] * ((-len(h)) % self.block) |
|
|
| def pick(self, paper, prefix): |
| p = self.encode(paper) |
| if len(p) > self.paper_tokens: |
| paper = self.tok.decode(p[:self.paper_tokens]) |
| ids = prefix + self.encode(tail_text(paper) + self.post) |
| return self.check_length(ids) |
|
|
| def yesno(self, paper, topic, definition=None): |
| p = self.encode(paper) |
| if len(p) > self.paper_tokens: |
| paper = self.tok.decode(p[:self.paper_tokens]) |
| text = self.tok.apply_chat_template(messages(yesno_prompt(paper, topic, definition)), tokenize=False, |
| add_generation_prompt=True, enable_thinking=False) |
| return self.check_length(self.encode(text)) |
|
|
| def check_length(self, ids): |
| if len(ids) > self.max_tokens: |
| raise ValueError(f"Prompt has {len(ids)} tokens; limit {self.max_tokens}. Shorten descriptions/labels, " |
| "reduce paper_tokens, or split the taxonomy. No options were silently truncated.") |
| return ids |
|
|
|
|
| class ClassifierTasks: |
| """Shared task API for the Transformers and vLLM backends.""" |
| def classify_many(self, papers, options, scheme=SCHEME, question=QUESTION, *, |
| describe=True, parents=True, none_option=False, orders=1, seed=0): |
| papers = list(papers) |
| if not papers: |
| return [] |
| if not isinstance(orders, int) or orders < 1: |
| raise ValueError("orders must be a positive integer") |
| opts = self.inputs.normalize_options(options, none_option) |
| rng = random.Random(seed);k = len(opts) |
| total = [[0.0] * k for _ in papers] |
| for order in range(orders): |
| perm = list(range(k)) if order == 0 else rng.sample(range(k), k) |
| prefix = self.inputs.prefix([opts[j] for j in perm], scheme, question, describe, parents) |
| rows = [self.inputs.pick(p, prefix) for p in papers] |
| probs = self._scores(rows, [self.cfg["code_ids"][:k]] * len(rows), shared_prefix=True) |
| for t, p in zip(total, probs): |
| for pos, j in enumerate(perm): |
| t[j] += p[pos] / orders |
| return [dict(sorted(((o["id"], p) for o, p in zip(opts, vals)), key=lambda x: -x[1])) for vals in total] |
|
|
| def classify(self, paper, options, scheme=SCHEME, question=QUESTION, **kwargs): |
| return self.classify_many([paper], options, scheme, question, **kwargs)[0] |
|
|
| def is_on_topic_many(self, papers, topic, definition=None): |
| rows = [self.inputs.yesno(p, topic, definition) for p in papers] |
| if not rows: |
| return [] |
| return [r[0] for r in self._scores(rows, [[self.cfg["yes"], self.cfg["no"]]] * len(rows))] |
|
|
| def is_on_topic(self, paper, topic, definition=None): |
| return self.is_on_topic_many([paper], topic, definition)[0] |
|
|
|
|
| class PaperClassifier(ClassifierTasks): |
| def __init__(self, model="maximerivest/paper-classifier-2B-v2", *, revision=None, device=None, |
| dtype=None, batch_size=4, paper_tokens=700): |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| dtype = dtype or (torch.bfloat16 if self.device.startswith("cuda") else torch.float32) |
| if batch_size < 1: |
| raise ValueError("batch_size must be positive") |
| self.batch_size = batch_size |
| self.cfg = read_config(model, revision) |
| self.tok = AutoTokenizer.from_pretrained(model, revision=revision) |
| self.tok.padding_side = "left" |
| self.inputs = PaperInputs(self.tok, self.cfg, paper_tokens=paper_tokens) |
| self.model = AutoModelForCausalLM.from_pretrained(model, revision=revision, dtype=dtype, |
| attn_implementation="sdpa").to(self.device).eval() |
|
|
| def _scores(self, input_ids, answer_ids, shared_prefix=False): |
| import torch |
| out = [] |
| with torch.inference_mode(): |
| for i in range(0, len(input_ids), self.batch_size): |
| enc = self.tok.pad({"input_ids": input_ids[i:i+self.batch_size]}, return_tensors="pt").to(self.device) |
| h = self.model.model(**enc, use_cache=False).last_hidden_state[:, -1, :].float() |
| W = self.model.get_output_embeddings().weight |
| for row, ids in zip(h, answer_ids[i:i+self.batch_size]): |
| out.append(torch.softmax(row @ W[ids].float().T, -1).cpu().tolist()) |
| return out |
|
|
|
|
| if __name__ == "__main__": |
| import sys |
| model = sys.argv[1] if len(sys.argv)>1 else "maximerivest/paper-classifier-2B-v2" |
| clf = PaperClassifier(model) |
| paper = paper_text("Habitat fragmentation and tropical forest birds", |
| "We surveyed bird communities across 40 forest fragments. Species richness declined with " |
| "fragment area and increased with connectivity.") |
| print(clf.classify(paper, ["Ecology", "Ornithology", "Particle physics", "Economics"], scheme="Fields of science")) |
| print("P(topic is central):", clf.is_on_topic(paper, "habitat fragmentation")) |
|
|
