{text}
" 176 | 177 | # return (LazyLines(stream) 178 | # .mutate(abstract=lambda d: d['text']) 179 | # .pipe(attach_docs, nlp=self.nlp, model=model) 180 | # .mutate(cats = lambda d: d['doc'].cats) 181 | # .keep(lambda d: d['cats'].get(tag, 0.0) > 0.6) 182 | # .mutate(html=lambda d: render_html(d['doc']), 183 | # n_sents=lambda d: len(d['doc'].spans["sc"]), 184 | # link=lambda d: d['meta']['url'])) 185 | 186 | # def build(self): 187 | # config = srsly.read_yaml(CONFIG_FILE) 188 | # for section in tqdm(config['sections'], desc="Looping over tags."): 189 | # print(section['name']) 190 | # section["content"] = self.fetch_tag_candidate_stream(tag=section['tag']).head(20).collect() 191 | # import pprint 192 | # pprint.pprint(section["content"]) 193 | 194 | # template = jinja2.Template(Path(TEMPLATE_PATH).read_text()) 195 | # Path("site.html").write_text(template.render(sections=config['sections'], today=dt.date.today())) 196 | 197 | 198 | # def evaluate(self): 199 | # from ._benchmark import benchmark 200 | # annotated, found_tags = self.fetch_annotated_data() 201 | # benchmark(annotated, tags=["new-dataset", "data-quality"]) 202 | 203 | # def push_wandb(self): 204 | # ... 205 | 206 | # def pull_wandb(self): 207 | # ... 208 | 209 | 210 | # if __name__ == "__main__": 211 | # Frontpage(config=srsly.read_yaml("config.yml")).build() 212 | -------------------------------------------------------------------------------- /frontpage/__main__.py: -------------------------------------------------------------------------------- 1 | import os 2 | import datetime as dt 3 | from pathlib import Path 4 | 5 | from jinja2 import Template 6 | from radicli import Radicli, Arg 7 | 8 | from .utils import console 9 | from .constants import TEMPLATE_PATH, TRAINED_FOLDER, SITE_PATH 10 | 11 | cli = Radicli() 12 | 13 | 14 | @cli.command("download") 15 | def download(): 16 | """Download new data.""" 17 | from .download import main as download_data 18 | download_data() 19 | 20 | 21 | @cli.command("index", 22 | kind=Arg(help="Can be lunr/simsity"), 23 | level=Arg(help="Can be sentence/abstract") 24 | ) 25 | def index_cli(kind:str, level:str): 26 | """Creates index for annotation.""" 27 | from .datastream import DataStream 28 | 29 | DataStream().create_index(level=level, kind=kind) 30 | 31 | 32 | @cli.command("preprocess") 33 | def preprocess_cli(): 34 | """Dedup and process data for faster processing.""" 35 | from .datastream import DataStream 36 | DataStream().save_clean_download_stream() 37 | 38 | 39 | @cli.command("annotate") 40 | def annotate(): 41 | """Annotate new examples.""" 42 | def run_questions(): 43 | import questionary 44 | from .constants import LABELS, DATA_LEVELS 45 | results = {} 46 | results["label"] = questionary.select( 47 | "Which label do you want to annotate?", 48 | choices=LABELS, 49 | ).ask() 50 | 51 | results["level"] = questionary.select( 52 | "What view of the data do you want to take?", 53 | choices=DATA_LEVELS, 54 | ).ask() 55 | 56 | if results["level"] == "abstract": 57 | choices = ["second-opinion", "search-engine", "simsity", "random"] 58 | else: 59 | choices = ["simsity", "search-engine", "active-learning", "random"] 60 | 61 | results["tactic"] = questionary.select( 62 | "Which tactic do you want to apply?", 63 | choices=choices, 64 | ).ask() 65 | 66 | results['setting'] = '' 67 | if results["tactic"] in ["simsity", "search-engine"]: 68 | results["setting"] = questionary.text( 69 | "What query would you like to use?", "" 70 | ).ask() 71 | 72 | if results["tactic"] == "active-learning": 73 | results["setting"] = questionary.select( 74 | "What should the active learning method prefer?", 75 | choices=["positive class", "uncertainty", "negative class"], 76 | ).ask() 77 | return results 78 | 79 | results = run_questions() 80 | from .recipe import annotate_prodigy 81 | annotate_prodigy(results) 82 | 83 | @cli.command("annotprep") 84 | def annotprep(): 85 | """Prepares data for training.""" 86 | from .datastream import DataStream 87 | DataStream().save_train_stream() 88 | 89 | 90 | @cli.command("train") 91 | def train(): 92 | """Trains a new model on the data.""" 93 | from .datastream import DataStream 94 | from .modelling import SentenceModel 95 | examples = DataStream().get_train_stream() 96 | SentenceModel().train(examples=examples).to_disk() 97 | 98 | 99 | @cli.command("pretrain") 100 | def pretrain(): 101 | """Trains a new featurizer, set-fit style.""" 102 | from .datastream import DataStream 103 | from .modelling import SentenceModel 104 | examples = DataStream().get_train_stream() 105 | SentenceModel().pretrain(examples=examples) 106 | 107 | 108 | @cli.command("stats") 109 | def stats(): 110 | """Show annotation stats""" 111 | from .datastream import DataStream 112 | DataStream().show_annot_stats() 113 | 114 | 115 | @cli.command( 116 | "build", 117 | retrain=Arg("--retrain", "-rt", help="Retrain model?"), 118 | prep=Arg("--preprocess", "-pr", help="Preprocess again?") 119 | ) 120 | def build(retrain: bool = False, prep:bool = False): 121 | """Build a new site""" 122 | from .datastream import DataStream 123 | if prep: 124 | preprocess_cli() 125 | if retrain: 126 | train() 127 | console.log("Starting site build process") 128 | sections = DataStream().get_site_content() 129 | template = Template(Path(TEMPLATE_PATH).read_text()) 130 | rendered = template.render(sections=sections, today=dt.date.today()) 131 | SITE_PATH.write_text(rendered) 132 | console.log("Site built.") 133 | 134 | 135 | @cli.command("artifact", 136 | action=Arg(help="Can be upload/download"), 137 | ) 138 | def artifact(action:str): 139 | """Upload/download from wandb""" 140 | import wandb 141 | from dotenv import load_dotenv 142 | from frontpage.constants import PRETRAINED_FOLDER 143 | load_dotenv() 144 | run = wandb.init(os.getenv("WANDB_API_KEY")) 145 | if action == "upload": 146 | artifact = wandb.Artifact(name='custom-sbert-emb', type="model") 147 | artifact.add_dir(local_path=PRETRAINED_FOLDER) 148 | run = wandb.init(project="arxiv-frontpage", job_type="upload") 149 | run.log_artifact(artifact) 150 | if action == "download": 151 | if not PRETRAINED_FOLDER.exists(): 152 | run = wandb.init(project="arxiv-frontpage", job_type="download") 153 | artifact = run.use_artifact('custom-sbert-emb:latest') 154 | console.log(f"Could not find {PRETRAINED_FOLDER}. So will download from wandb.") 155 | artifact.download(PRETRAINED_FOLDER) 156 | else: 157 | console.log(f"{PRETRAINED_FOLDER} already exists. Skip wandb download.") 158 | 159 | 160 | @cli.command("search") 161 | def search(): 162 | """Annotate new examples.""" 163 | import questionary 164 | from simsity import load_index 165 | from .modelling import SentenceModel 166 | enc = SentenceModel().encoder 167 | index = load_index("indices/simsity/sentence", encoder=enc) 168 | while True: 169 | query = questionary.text("Query:").ask() 170 | texts, dists = index.query([query], n=5) 171 | for t in texts: 172 | print(t) 173 | 174 | if __name__ == "__main__": 175 | cli.run() 176 | -------------------------------------------------------------------------------- /frontpage/_benchmark.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | import itertools as it 3 | 4 | import tqdm 5 | import srsly 6 | import numpy as np 7 | import polars as pl 8 | from dotenv import load_dotenv 9 | from sklearn.svm import SVC 10 | from sklearn.pipeline import make_pipeline 11 | from sklearn.linear_model import LogisticRegression 12 | from sklearn.feature_extraction.text import HashingVectorizer 13 | from sklearn.metrics import classification_report 14 | from sklearn.model_selection import StratifiedKFold, train_test_split 15 | from embetter.text import SentenceEncoder, spaCyEncoder 16 | from embetter.external import CohereEncoder, OpenAIEncoder 17 | from embetter.utils import cached 18 | from sklearn.pipeline import make_pipeline, make_union 19 | from sklearn.decomposition import TruncatedSVD 20 | from embetter.finetune import ForwardFinetuner, ContrastiveFinetuner 21 | from sklearn.preprocessing import FunctionTransformer 22 | 23 | from frontpage.datastream import DataStream 24 | 25 | 26 | load_dotenv() 27 | 28 | def grid(**kwargs): 29 | res = [{k: v for k, v in zip(kwargs.keys(), prod)} 30 | for prod in it.product(*[v for v in kwargs.values()])] 31 | return tqdm.tqdm(res) 32 | 33 | 34 | datastream = DataStream() 35 | 36 | k_folder = StratifiedKFold(n_splits=10) 37 | 38 | encoders = { 39 | "spacy": spaCyEncoder("en_core_web_md"), 40 | "sbert": SentenceEncoder(), 41 | "hash_lg": HashingVectorizer(), 42 | "hash_sm": HashingVectorizer(n_features=2**14), 43 | "openai": OpenAIEncoder(), 44 | "cohere": CohereEncoder(), 45 | } 46 | 47 | encoders["multi"] = make_union( 48 | encoders["sbert"], 49 | make_pipeline( 50 | HashingVectorizer(n_features=10_000), 51 | TruncatedSVD(), 52 | ) 53 | ) 54 | 55 | tuners = { 56 | "forward": lambda: ForwardFinetuner(hidden_dim=300), 57 | "contrast": lambda: ContrastiveFinetuner(hidden_dim=300), 58 | "none": lambda: FunctionTransformer() 59 | } 60 | 61 | for name, enc in encoders.items(): 62 | if name not in ["multi", "hash_lg", "hash_sm"]: 63 | encoders[name] = cached(f"cache/{str(type(enc))}", enc) 64 | 65 | models = { 66 | "logistic": LogisticRegression(class_weight="balanced", max_iter=1000), 67 | "svm": SVC(class_weight="balanced") 68 | } 69 | 70 | def calc_stats(pred_valid, y_valid): 71 | return {**classification_report(pred_valid, y_valid, output_dict=True)['1'], "accuracy": float(np.mean(pred_valid == y_valid))} 72 | 73 | 74 | def run_benchmark_k_fold(label, model, encoder, tuner): 75 | res = {"label": label, "model": model, "encoder": encoder, "tuner": tuner, "method": "k_fold"} 76 | pipe = make_pipeline(encoders[encoder], tuners[tuner](), models[model]) 77 | examples = datastream.get_train_stream() 78 | X = [ex['text'] for ex in examples if label in ex['cats']] 79 | y = [ex['cats'][label] for ex in examples if label in ex['cats']] 80 | folds = k_folder.split(X, y) 81 | for i, (train_idx, valid_idx) in enumerate(folds): 82 | X_train = [str(x) for x in np.array(X)[train_idx]] 83 | X_valid = [str(x) for x in np.array(X)[valid_idx]] 84 | y_train = np.array(y)[train_idx] 85 | y_valid = np.array(y)[valid_idx] 86 | pipe.fit(X_train, y_train) 87 | valid_pred = pipe.predict(X_valid) 88 | stats = calc_stats(valid_pred, y_valid) 89 | res = {**res, **stats, "data_size": len(y), "i": i} 90 | yield res 91 | 92 | 93 | def run_benchmark_train_size(label, model, encoder, tuner): 94 | res = {"label": label, "model": model, "encoder": encoder, "tuner": tuner, "method": "train_size"} 95 | pipe = make_pipeline(encoders[encoder], tuners[tuner](), models[model]) 96 | examples = datastream.get_train_stream() 97 | X = [ex['text'] for ex in examples if label in ex['cats']] 98 | y = [ex['cats'][label] for ex in examples if label in ex['cats']] 99 | X_train, X_valid, y_train, y_valid = train_test_split(X, y, test_size=0.2) 100 | for p in [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]: 101 | idx = int(len(X_train) * p) 102 | X_train_use = [str(x) for x in np.array(X_train)[:idx]] 103 | y_train_use = np.array(y_train)[:idx] 104 | pipe.fit(X_train_use, y_train_use) 105 | valid_pred = pipe.predict(X_valid) 106 | stats = calc_stats(valid_pred, y_valid) 107 | res = {**res, **stats, "data_size": len(y), "p": p} 108 | yield res 109 | 110 | 111 | if __name__ == "__main__": 112 | settings = grid( 113 | label=["new-dataset"], 114 | encoder=["sbert", "openai", "cohere", "multi"], 115 | model=["logistic", "svm"], 116 | tuner=["contrast", "forward", "none"] 117 | ) 118 | 119 | stats = (ex for setting in settings for ex in run_benchmark_k_fold(**setting)) 120 | 121 | if Path("benchmark_kfold.jsonl").exists(): 122 | Path("benchmark_kfold.jsonl").unlink() 123 | srsly.write_jsonl("benchmark_kfold.jsonl", stats) 124 | 125 | stats = (ex for setting in settings for ex in run_benchmark_train_size(**setting)) 126 | 127 | if Path("benchmark_train_size.jsonl").exists(): 128 | Path("benchmark_train_size.jsonl").unlink() 129 | srsly.write_jsonl("benchmark_train_size.jsonl", stats) 130 | 131 | pl.Config.set_tbl_rows(100) 132 | pl.Config.set_tbl_width_chars(1000) 133 | 134 | # print( 135 | # pl.read_ndjson("benchmark.jsonl") 136 | # .groupby("label","model","encoder","tuner") 137 | # .agg( 138 | # pl.mean("recall"), 139 | # pl.mean("precision"), 140 | # pl.mean("f1-score"), 141 | # pl.mean("accuracy"), 142 | # ).sort("f1-score") 143 | # ) 144 | 145 | -------------------------------------------------------------------------------- /frontpage/constants.py: -------------------------------------------------------------------------------- 1 | import srsly 2 | from pathlib import Path 3 | from typing import Literal 4 | 5 | from .types import Config 6 | 7 | # Paths and folders 8 | DATA_FOLDER = Path("data") 9 | ANNOT_FOLDER = DATA_FOLDER / "annot" 10 | INDICES_FOLDER = Path("indices") 11 | CLEAN_DOWNLOADS_FOLDER = DATA_FOLDER / Path("cleaned") 12 | DOWNLOADS_FOLDER = DATA_FOLDER / "downloads" 13 | ANNOT_PATH = ANNOT_FOLDER / "annotations.jsonl" 14 | ACTIVE_LEARN_PATH = ANNOT_FOLDER / "active-learn.jsonl" 15 | SECOND_OPINION_PATH = ANNOT_FOLDER / "second-opinion.jsonl" 16 | TRAINED_FOLDER = Path("training") 17 | TEMPLATE_PATH = Path("templates/home.html") 18 | CONFIG_FILE = "config.yml" 19 | PRETRAINED_FOLDER = TRAINED_FOLDER / "custom-sbert-emb" 20 | 21 | # Cache paths 22 | EMBETTER_CACHE = Path("cache") / "embetter" 23 | 24 | # Possible values 25 | DATA_LEVELS = ["sentence", "abstract"] 26 | DATA_LEVELS_TYPE = Literal["sentence", "abstract"] 27 | CONFIG = Config(**srsly.read_yaml(CONFIG_FILE)) 28 | LABELS = [s.label for s in CONFIG.sections] 29 | THRESHOLDS = {s.label: s.threshold for s in CONFIG.sections} 30 | SITE_PATH = Path("index.html") 31 | -------------------------------------------------------------------------------- /frontpage/download.py: -------------------------------------------------------------------------------- 1 | import datetime as dt 2 | from pathlib import Path 3 | from typing import List 4 | 5 | import srsly 6 | import tqdm 7 | import arxiv 8 | from arxiv import Result 9 | from retry import retry 10 | import spacy 11 | from spacy.language import Language 12 | from .types import ArxivArticle 13 | from rich.console import Console 14 | 15 | console = Console() 16 | 17 | 18 | def age_in_days(res: Result) -> float: 19 | """Get total seconds from now from Arxiv result""" 20 | now = dt.datetime.now(dt.timezone.utc) 21 | return (now - res.published).total_seconds() / 3600 / 24 22 | 23 | 24 | def parse(res: Result, nlp: Language) -> ArxivArticle: 25 | """Parse proper Pydantic object from Arxiv""" 26 | summary = res.summary.replace("\n", " ") 27 | doc = nlp(summary) 28 | sents = [s.text for s in doc.sents] 29 | 30 | return ArxivArticle( 31 | created=str(res.published)[:19], 32 | title=str(res.title), 33 | abstract=summary, 34 | sentences=sents, 35 | url=str(res.entry_id) 36 | ) 37 | 38 | @retry(tries=5, delay=1, backoff=2) 39 | def main(): 40 | nlp = spacy.load("en_core_web_sm", disable=["ner", "lemmatizer", "tagger"]) 41 | console.log(f"Starting arxiv search.") 42 | items = arxiv.Search( 43 | query="and", 44 | max_results=200, 45 | sort_by=arxiv.SortCriterion.SubmittedDate, 46 | ) 47 | 48 | results = list(items.results()) 49 | 50 | console.log(f"Found {len(results)} results.") 51 | 52 | articles = [dict(parse(r, nlp=nlp)) 53 | for r in tqdm.tqdm(results) 54 | if age_in_days(r) < 2.5 and r.primary_category.startswith("cs")] 55 | 56 | dist = [age_in_days(r) for r in results] 57 | if dist: 58 | console.log(f"Minimum article age: {min(dist)}") 59 | console.log(f"Maximum article age: {max(dist)}") 60 | articles_dict = {ex['title']: ex for ex in articles} 61 | most_recent = list(sorted(Path("data/downloads/").glob("*.jsonl")))[-1] 62 | old_articles_dict = {ex['title']: ex for ex in srsly.read_jsonl(most_recent)} 63 | 64 | new_articles = [ex for title, ex in articles_dict.items() if title not in old_articles_dict.keys()] 65 | old_articles = [ex for title, ex in articles_dict.items() if title in old_articles_dict.keys()] 66 | if old_articles: 67 | console.log(f"Found {len(old_articles)} old articles in current batch. Skipping.") 68 | if new_articles: 69 | console.log(f"Found {len(new_articles)} new articles in current batch to write.") 70 | filename = str(dt.datetime.now()).replace(" ", "-")[:13] + "h.jsonl" 71 | srsly.write_jsonl(Path("data") / "downloads" / filename, new_articles) 72 | console.log(f"Wrote {len(new_articles)} articles into {filename}.") 73 | -------------------------------------------------------------------------------- /frontpage/modelling.py: -------------------------------------------------------------------------------- 1 | from typing import List, Dict 2 | from pathlib import Path 3 | from functools import cached_property 4 | 5 | from wasabi import Printer 6 | import numpy as np 7 | from skops.io import dump, load 8 | from embetter.utils import cached 9 | from sklearn.linear_model import LogisticRegression 10 | 11 | from .constants import TRAINED_FOLDER, LABELS, EMBETTER_CACHE, PRETRAINED_FOLDER 12 | from .utils import console 13 | 14 | msg = Printer() 15 | 16 | 17 | class SentenceModel: 18 | def __init__(self, labels=LABELS) -> None: 19 | self.labels = labels 20 | self._models = {k: LogisticRegression(class_weight="balanced") for k in self.labels} 21 | 22 | def train(self, examples): 23 | X = self.featurizer.transform([ex["text"] for ex in examples]) 24 | for task, model in self._models.items(): 25 | xs = np.array([X[i] for i, ex in enumerate(examples) if task in ex['cats']]) 26 | ys = np.array( 27 | [ex['cats'][task] for ex in examples if task in ex['cats']], dtype=int 28 | ) 29 | model.fit(xs, ys) 30 | console.log(f"Trained the [bold]{task}[/bold] task, using {len(xs)} examples.") 31 | return self 32 | 33 | def __call__(self, text:str) -> Dict: 34 | result = {} 35 | X = self.featurizer.transform([text]) 36 | for label in self.labels: 37 | proba = self._models[label].predict_proba(X)[0, 1] 38 | result[label] = float(proba) 39 | return result 40 | 41 | def predict(self, texts: List[str]) -> List[Dict]: 42 | X = self.featurizer.transform(texts) 43 | result = [{} for _ in texts] 44 | for label in self.labels: 45 | probas = self._models[label].predict_proba(X)[:, 1] 46 | for i, proba in enumerate(probas): 47 | result[i][label] = float(proba) 48 | return result 49 | 50 | @cached_property 51 | def encoder(self): 52 | from embetter.text import SentenceEncoder 53 | encoder = SentenceEncoder() 54 | encoder = cached(EMBETTER_CACHE / "sbert", encoder) 55 | return encoder 56 | 57 | @cached_property 58 | def featurizer(self): 59 | from embetter.text import SentenceEncoder 60 | if not Path(PRETRAINED_FOLDER).exists(): 61 | console.log("Did not find pretrained model. Falling back.") 62 | return self.encoder 63 | console.log(f"Will use custom model found in {PRETRAINED_FOLDER}") 64 | return SentenceEncoder(PRETRAINED_FOLDER) 65 | 66 | def pretrain(self, examples): 67 | from sentence_transformers import SentenceTransformer, InputExample, losses 68 | from torch.utils.data import DataLoader 69 | from embetter.finetune._contrastive import generate_pairs_batch 70 | 71 | console.log("Starting pretraining sequence.") 72 | all_pairs = [] 73 | for label in LABELS: 74 | subset = [ex for ex in examples if label in ex['cats']] 75 | pairs = generate_pairs_batch([ex['cats'][label] for ex in subset], n_neg=1) 76 | all_pairs.extend(pairs) 77 | 78 | input_examples = [] 79 | for pair in pairs: 80 | text1 = examples[pair.i1]['text'] 81 | text2 = examples[pair.i2]['text'] 82 | input_examples.append(InputExample(texts=[text1, text2], label=float(pair.label))) 83 | model = SentenceTransformer('all-MiniLM-L6-v2') 84 | 85 | train_dataloader = DataLoader(input_examples, shuffle=True, batch_size=16) 86 | train_loss = losses.CosineSimilarityLoss(model) 87 | 88 | console.log("Pairs generated. About to tune model.") 89 | model.fit( 90 | train_objectives=[(train_dataloader, train_loss)], 91 | epochs=3, 92 | warmup_steps=100, 93 | output_path=str(PRETRAINED_FOLDER) 94 | ) 95 | console.log(f"New encoder saved at {PRETRAINED_FOLDER}") 96 | 97 | @cached_property 98 | def nlp(self): 99 | import spacy 100 | return spacy.load("en_core_web_sm", disable=["ner", "lemmatizer", "tagger"]) 101 | 102 | def to_disk(self, path: Path=TRAINED_FOLDER): 103 | if not Path(path).exists(): 104 | Path(path).mkdir(exist_ok=True, parents=True) 105 | # Delete old files 106 | if Path(path).exists(): 107 | for p in Path(path).glob("*.h5"): 108 | p.unlink() 109 | # Write new files 110 | for name, clf in self._models.items(): 111 | dump(clf, Path(path) / f"{name}.h5") 112 | console.log(f"Model saved in folder: [bold]{path}[/bold].") 113 | 114 | @classmethod 115 | def from_disk(cls, path: Path=TRAINED_FOLDER): 116 | if not Path(path).exists(): 117 | raise RuntimeError("You need to train a model beforehand.") 118 | models = {} 119 | for f in Path(path).glob("*.h5"): 120 | models[f.stem] = load(f, trusted=True) 121 | 122 | model = SentenceModel(labels=models.keys()) 123 | model._models = models 124 | console.log(f"Model loaded from: [bold]{path}[/bold].") 125 | return model 126 | -------------------------------------------------------------------------------- /frontpage/recipe.py: -------------------------------------------------------------------------------- 1 | import spacy 2 | 3 | import prodigy 4 | 5 | from .datastream import DataStream 6 | from .utils import console 7 | datastream = DataStream() 8 | 9 | 10 | 11 | @prodigy.recipe("textcat.arxiv.sentence", 12 | dataset=("The dataset to save", "positional", None, str), 13 | label=("The label to annotate", "positional", None, str), 14 | tactic=("The tactic to retreive relevant examples", "positional", None, str), 15 | setting=("Additional setting for the tactic", "positional", None, str), 16 | ) 17 | def arxiv_sentence(dataset, label, tactic, setting): 18 | """Very general recipe to annotate sentences, using different ordering techniques.""" 19 | from prodigy import set_hashes 20 | if tactic == "simsity": 21 | console.log("Setting up simsity stream") 22 | stream = datastream.get_ann_stream(query=setting, level="sentence") 23 | elif tactic == "random": 24 | console.log("Setting up randomized stream") 25 | stream = datastream.get_random_stream(level="sentence") 26 | elif tactic == "active-learning": 27 | console.log("Setting up active learning") 28 | stream = datastream.get_active_learn_stream(label=label, preference=setting) 29 | elif tactic == "search-engine": 30 | console.log("Setting up lunr query") 31 | stream = datastream.get_lunr_stream(query=setting, level="sentence") 32 | else: 33 | raise ValueError("This should never happen.") 34 | 35 | return { 36 | "dataset": dataset, 37 | "stream": (set_hashes({**ex, "label": label}) for ex in stream), 38 | "view_id": "classification", 39 | "config":{ 40 | "exclude_by": "input" 41 | } 42 | } 43 | 44 | 45 | @prodigy.recipe("textcat.arxiv.abstract", 46 | dataset=("The dataset to save", "positional", None, str), 47 | label=("The label to annotate", "positional", None, str), 48 | tactic=("The tactic to retreive relevant examples", "positional", None, str), 49 | setting=("Additional setting for the tactic", "positional", None, str), 50 | ) 51 | def arxiv_abstract(dataset, label, tactic, setting): 52 | """Very general recipe to annotate sentences, using different ordering techniques.""" 53 | from prodigy.components.preprocess import add_tokens 54 | from prodigy import set_hashes 55 | 56 | if tactic == "simsity": 57 | console.log("Setting up simsity stream") 58 | stream = datastream.get_ann_stream(query=setting, level="abstract") 59 | elif tactic == "random": 60 | console.log("Setting up randomized stream") 61 | stream = datastream.get_random_stream(level="abstract") 62 | elif tactic == "search-engine": 63 | console.log("Setting up lunr query") 64 | stream = datastream.get_lunr_stream(query=setting, level="abstract") 65 | elif tactic == "second-opinion": 66 | console.log("Setting up second opinion") 67 | stream = datastream.get_second_opinion_stream(label=label, min_sents=1, max_sents=2) 68 | else: 69 | raise ValueError("This should never happen.") 70 | 71 | nlp = spacy.blank("en") 72 | stream = ({**ex, "label": label} for ex in stream) 73 | stream = add_tokens(nlp, stream) 74 | return { 75 | "dataset": dataset, 76 | "stream": (set_hashes(ex) for ex in stream), 77 | "view_id": "blocks", 78 | "config": { 79 | "labels": [label], 80 | "blocks": [ 81 | {"view_id": "ner_manual"}, 82 | ], 83 | "exclude_by": "input" 84 | } 85 | } 86 | 87 | 88 | def annotate_prodigy(results): 89 | from prodigy.app import server 90 | from prodigy.core import Controller 91 | 92 | dataset_name = datastream.get_dataset_name(results['label'], results['level']) 93 | name = "textcat.arxiv.sentence" if results['level'] == 'sentence' else "textcat.arxiv.abstract" 94 | if results['level'] == 'sentence': 95 | ctrl_data = arxiv_sentence(dataset_name, results['label'], results['tactic'], results['setting']) 96 | else: 97 | ctrl_data = arxiv_abstract(dataset_name, results['label'], results['tactic'], results['setting']) 98 | controller = Controller.from_components(name, ctrl_data) 99 | server(controller, controller.config) 100 | -------------------------------------------------------------------------------- /frontpage/types.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | from pydantic import BaseModel 3 | 4 | 5 | class ArxivArticle(BaseModel): 6 | created: str 7 | title: str 8 | abstract: str 9 | sentences: List[str] 10 | url: str 11 | 12 | 13 | class LabelConfig(BaseModel): 14 | name: str 15 | instructions: str 16 | label: str 17 | threshold: float = 0.7 18 | 19 | 20 | class Config(BaseModel): 21 | n_per_section: int = 20 22 | sections: List[LabelConfig] 23 | -------------------------------------------------------------------------------- /frontpage/utils.py: -------------------------------------------------------------------------------- 1 | from rich.console import Console 2 | import itertools as it 3 | 4 | from spacy.tokens import Span 5 | 6 | 7 | console = Console() 8 | 9 | 10 | def batched(iterable, n=56): 11 | "Batch data into tuples of length n. The last batch may be shorter." 12 | if n < 1: 13 | raise ValueError('n must be at least one') 14 | iters = iter(iterable) 15 | while batch := tuple(it.islice(iters, n)): 16 | yield batch 17 | 18 | 19 | def dedup_stream(stream, key="text"): 20 | uniq = {} 21 | for ex in stream: 22 | uniq[hash(ex[key])] = ex 23 | for ex in uniq.values(): 24 | yield ex 25 | 26 | 27 | def add_rownum(stream): 28 | for i, ex in enumerate(stream): 29 | yield {"text": ex["text"], "idx": i} 30 | 31 | 32 | def attach_docs(lines, nlp, label): 33 | tuples = ((eg['text'], eg) for eg in lines) 34 | for doc, eg in nlp.pipe(tuples, as_tuples=True): 35 | eg['doc'] = sentence_classifier(doc, eg['preds'], label) 36 | yield eg 37 | 38 | 39 | def sentence_classifier(doc, preds, label): 40 | doc.spans["sc"] = [] 41 | for sent, pred in zip(doc.sents, preds): 42 | for k, p in pred.items(): 43 | if p >= 0.6: 44 | if k == label: 45 | doc.spans["sc"].append(Span(doc, sent.start, sent.end, k)) 46 | doc.cats[k] = max(doc.cats.get(k, 0.0), p) 47 | return doc 48 | 49 | 50 | def attach_spans(stream, label, min_spans=1, max_spans=1): 51 | for ex in stream: 52 | spans = [] 53 | for spansvals in ex['doc'].spans.values(): 54 | for span in spansvals: 55 | spans.append( 56 | { 57 | "token_start": span.start, 58 | "token_end": span.end - 1, 59 | "start": span.start_char, 60 | "end": span.end_char, 61 | "text": span.text, 62 | "label": label, 63 | } 64 | ) 65 | ex["spans"] = spans 66 | del ex["doc"] 67 | if len(spans) >= min_spans: 68 | if len(spans) <= max_spans: 69 | yield ex 70 | 71 | 72 | def add_predictions(stream, model): 73 | for ex in stream: 74 | preds = model.predict(ex['sentences']) 75 | ex['preds'] = preds 76 | ex['created'] = ex['created'][:10] 77 | yield ex 78 | 79 | 80 | def _abstract_single_annot_to_sent(example, nlp, label): 81 | """Takens an annotation from abstract level and turns it into a training example""" 82 | text = example['text'] 83 | if example['answer'] == "accept" and "spans" in example: 84 | for span in example['spans']: 85 | yield {"text": text[span['start']: span['end']], label: 1} 86 | for span in example['spans']: 87 | text = text.replace(text[span['start']: span['end']], "") 88 | for sent in nlp(text).sents: 89 | if len(sent.text) > 5: 90 | yield {"text": sent.text, label: 0} 91 | 92 | 93 | def abstract_annot_to_sent(examples, nlp, label): 94 | """Takens an annotation from abstract level and turns it into a training example""" 95 | for ex in examples: 96 | for annot in _abstract_single_annot_to_sent(ex, nlp, label): 97 | yield annot 98 | -------------------------------------------------------------------------------- /images/active-teaching.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/koaning/arxiv-frontpage/260eaf2a63af68045367f0579161b6b414fc97ac/images/active-teaching.png -------------------------------------------------------------------------------- /images/multiheads.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/koaning/arxiv-frontpage/260eaf2a63af68045367f0579161b6b414fc97ac/images/multiheads.png -------------------------------------------------------------------------------- /images/sentence-model.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/koaning/arxiv-frontpage/260eaf2a63af68045367f0579161b6b414fc97ac/images/sentence-model.png -------------------------------------------------------------------------------- /prodigy.json: -------------------------------------------------------------------------------- 1 | { 2 | "exclude_by": "input" 3 | } -------------------------------------------------------------------------------- /requirements-build.txt: -------------------------------------------------------------------------------- 1 | srsly==2.4.6 2 | tqdm==4.65.0 3 | arxiv==1.4.2 4 | httpx==0.23.3 5 | spacy==3.5.0 6 | retry==0.9.2 7 | wandb==0.14.2 8 | embetter==0.5.0 9 | sentence-transformers==2.2.2 10 | simsity==0.5.5 11 | questionary==1.10.0 12 | radicli==0.0.25 13 | lazylines==0.0.4 14 | rich==13.4.2 15 | lunr==0.6.2 16 | skops==0.6.0 17 | python-dotenv==1.0.0 18 | scikit-learn==1.3 19 | numpy==1.26.4 20 | huggingface_hub==0.24.7 21 | -------------------------------------------------------------------------------- /requirements-download.txt: -------------------------------------------------------------------------------- 1 | srsly==2.4.6 2 | tqdm==4.65.0 3 | arxiv==1.4.2 4 | httpx==0.23.3 5 | spacy==3.5.0 6 | retry==0.9.2 7 | rich==13.4.2 8 | radicli==0.0.25 9 | numpy==1.26 10 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | 4 | setup( 5 | name="frontpage", 6 | version="0.1.0", 7 | description="Your frontpage.", 8 | packages=find_packages(exclude=["notebooks", "tests"]), 9 | ) 10 | -------------------------------------------------------------------------------- /taskfile.yml: -------------------------------------------------------------------------------- 1 | version: '3' 2 | 3 | tasks: 4 | index: 5 | deps: [lunr-sentence, lunr-abstract] 6 | desc: Creates the simsity/lunr indices. 7 | sources: 8 | - data/cleaned/*.jsonl 9 | cmds: 10 | - python -m frontpage index simsity sentence 11 | - python -m frontpage index simsity abstract 12 | 13 | lunr-sentence: 14 | desc: Creates the sentence-level lunr index. 15 | cmds: 16 | - python -m frontpage index lunr sentence 17 | sources: 18 | - data/cleaned/*.jsonl 19 | generates: 20 | - indices/lunr/sentence.json 21 | 22 | lunr-abstract: 23 | desc: Creates the abstract-level lunr index. 24 | cmds: 25 | - python -m frontpage index lunr abstract 26 | sources: 27 | - data/cleaned/*.jsonl 28 | generates: 29 | - indices/lunr/abstract.json 30 | 31 | clean: 32 | desc: Cleans the downloaded data. 33 | cmds: 34 | - python -c "from frontpage.datastream import DataStream; DataStream().save_clean_download_stream()" 35 | sources: 36 | - data/downloads/*.jsonl 37 | 38 | build-learned-streams: 39 | desc: Generates files for the active-learn-y annotation recipes. 40 | cmds: 41 | - python -c "from frontpage.datastream import DataStream; DataStream().build_active_learn_stream()" 42 | - python -c "from frontpage.datastream import DataStream; DataStream().build_second_opinion_stream()" 43 | sources: 44 | - data/cleaned/*.jsonl 45 | - training 46 | 47 | stats-out: 48 | desc: Runs `prodigy stats`, which could indicate that there are new annotations. 49 | cmds: 50 | - python -m prodigy stats -nf > /tmp/stats.jsonl 51 | generates: 52 | - tmp/stats.jsonl 53 | 54 | prepare-annotations: 55 | desc: Export the annotations from Prodigy so ML can be trained on it. 56 | cmds: 57 | - python -m frontpage annotprep 58 | sources: 59 | - /tmp/stats.jsonl 60 | generates: 61 | - data/annot/annotations.jsonl 62 | 63 | train: 64 | desc: Trains new classification heads on top of pretrained SBERT layer. 65 | cmds: 66 | - python -m frontpage train 67 | sources: 68 | - data/annot/annotations.jsonl 69 | - training/custom-sbert-emb 70 | 71 | pretrain: 72 | desc: Trains new SBERT representations. Can be expensive. Will also upload to wandb. 73 | cmds: 74 | - python -m frontpage pretrain 75 | - python -m frontpage artifact upload 76 | 77 | prepare: 78 | desc: Runs all the steps required to update the streams for annotation. 79 | cmds: 80 | - task: clean 81 | - task: index 82 | - task: stats-out 83 | - task: prepare-annotations 84 | - task: train 85 | - task: build-learned-streams 86 | 87 | build: 88 | desc: Constuct a new frontpage. Always retrains heads just in case. 89 | cmds: 90 | - task: clean 91 | # - python -m frontpage artifact download 92 | - task: train 93 | - python -m frontpage build 94 | - python -m http.server 95 | -------------------------------------------------------------------------------- /templates/home.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 || 19 | |
20 | Vincent's Arxiv FrontPage21 |22 | Generated on {{today}}. 23 | This frontpage is made by scraping arxiv and by running a sentence-model that detects if the abstract describes a paper about a topic of interest. One cool feature: it all pretty much runs via Github Actions. 24 |25 | |
26 |
| 30 | |
31 | {{section['name']}}32 | |
33 |
|
37 | {{content['created']}} 38 | |
39 |
40 |
41 |
42 | {{content['title']}}
43 |
44 |
60 |
45 |
59 |
46 | {%- for key in content['cats'] -%}
47 | cls:{{key}}: 0.{{(content['cats'][key] * 1000) | int | round}}
48 | {%- endfor -%}
49 |
50 | 51 | {{content['html']}} 52 | 53 |54 | 55 | link 56 | 57 | 58 | |
61 |