| """Classify papers as surveys or research papers with SciGenre. |
| |
| python predict.py papers.csv out.csv [--model REPO_ID_OR_PATH] |
| |
| papers.csv needs the columns `title` and `abstract`. The output adds `p_survey` and `category` (survey, research, |
| or no abstract when the abstract has fewer than ten words, which the model does not classify). |
| |
| The model input is model_text(title, abstract): the title and the abstract with markup, LaTeX and publisher |
| boilerplate removed, exactly as in training. The model never sees the venue. |
| """ |
| import argparse |
| import html |
| import re |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification |
|
|
| THRESHOLD, MAX_LEN, MIN_WORDS, SURVEY = 0.5, 384, 10, 0 |
|
|
|
|
| def clean_text(text) -> str: |
| """Remove markup, LaTeX and formatting artifacts; collapse whitespace.""" |
| if text is None or (isinstance(text, float) and text != text): |
| return "" |
| text = str(text) |
| text = re.sub(r"<tex-math.*?</tex-math>", " ", text, flags=re.DOTALL) |
| text = re.sub(r"<[^>]+>", " ", text) |
| text = re.sub(r"\$[^$]*\$", " ", text) |
| text = re.sub(r"\\[a-zA-Z]+\*?(?:\{[^}]*\})?", " ", text) |
| text = re.sub(r"[{}\[\]]", " ", text) |
| text = re.sub(r"^\s*abstract[:.\s-]*", "", text, flags=re.IGNORECASE) |
| return re.sub(r"\s+", " ", text).strip() |
|
|
|
|
| |
| |
| BOILERPLATE = re.compile( |
| r"(?:©|\(c\)\s*(?:19|20)\d\d|\bcopyright\s*(?:©|\(c\)|(?:19|20)\d\d)|\ball rights reserved\b|" |
| r"\bthis article is categorized under\b|\bexpected final online publication date\b|\bother information published in\b|" |
| r"\blicen[cs]e:\s*https?://|\bpublished online in wiley\b|\bcreative commons attribution\b)", |
| re.IGNORECASE, |
| ) |
|
|
|
|
| |
| SECTION_HEADER = re.compile( |
| r"(?:^|(?<=[.!?] ))(?:background|objectives?|purpose|aims?|methods?|methodology|design(?:/methodology/approach)?|" |
| r"approach|results?|findings|conclusions?|originality(?:/value)?|value|significance|context|introduction|" |
| r"motivation|limitations|implications)\s*:\s*", |
| re.IGNORECASE, |
| ) |
|
|
|
|
| def strip_boilerplate(text: str) -> str: |
| """Decode HTML entities, drop structured-abstract headings, then cut trailing boilerplate (marker in the last |
| 60% of the text) or drop the sentence with the marker.""" |
| text = SECTION_HEADER.sub("", html.unescape(text)) |
| text = re.sub(r"\bdoi:\s*10\.\S+", " ", text, flags=re.IGNORECASE) |
| m = BOILERPLATE.search(text) |
| while m: |
| if m.start() >= 0.4 * len(text): |
| text = text[:m.start()] |
| else: |
| s = text.rfind(". ", 0, m.start()) + 1 |
| e = text.find(". ", m.end()); e = len(text) if e < 0 else e + 1 |
| text = text[:s] + text[e:] |
| m = BOILERPLATE.search(text) |
| return re.sub(r"\s+", " ", text).strip() |
|
|
|
|
| def model_text(title, abstract) -> str: |
| """Input of the final model: title and abstract only, boilerplate removed (never the venue).""" |
| return f"{clean_text(title)}. {strip_boilerplate(clean_text(abstract))}" |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("papers"); ap.add_argument("out"); ap.add_argument("--model", default=".") |
| ap.add_argument("--batch", type=int, default=32) |
| a = ap.parse_args() |
| df = pd.read_csv(a.papers) |
| tok = AutoTokenizer.from_pretrained(a.model) |
| model = AutoModelForSequenceClassification.from_pretrained(a.model, dtype=torch.float32).eval() |
| ok = df["abstract"].fillna("").map(lambda s: len(clean_text(s).split()) >= MIN_WORDS).values |
| texts = [model_text(t, s) for t, s in zip(df["title"].fillna(""), df["abstract"].fillna(""))] |
| p = np.full(len(df), np.nan) |
| idx = np.where(ok)[0] |
| with torch.no_grad(): |
| for i in range(0, len(idx), a.batch): |
| b = idx[i:i + a.batch] |
| enc = tok([texts[j] for j in b], truncation=True, max_length=MAX_LEN, padding=True, return_tensors="pt") |
| p[b] = torch.softmax(model(**enc).logits.float(), -1)[:, SURVEY].numpy() |
| df["p_survey"] = p |
| df["category"] = np.where(~ok, "no abstract", np.where(p >= THRESHOLD, "survey", "research")) |
| df.to_csv(a.out, index=False) |
| print(df["category"].value_counts().to_string()) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
