nielsr/community-science-progress

1import dataclasses2from multiprocessing import cpu_count3import tqdm4import requests5import streamlit as st6 7import pandas as pd8from datasets import Dataset, load_dataset9from paperswithcode import PapersWithCodeClient10 11 12@dataclasses.dataclass(frozen=True)13class PaperInfo:14    date: str15    arxiv_id: str16    github: str17    title: str18    paper_page: str19    upvotes: int20    num_comments: int21 22 23def get_df(start_date: str = None, end_date: str = None) -> pd.DataFrame:24    """25    Load the initial dataset as a Pandas dataframe.26 27    One can optionally specify a start_date and end_date to only include data between these dates.28    """29 30    df = pd.merge(31        left=load_dataset("hysts-bot-data/daily-papers", split="train").to_pandas(),32        right=load_dataset("hysts-bot-data/daily-papers-stats", split="train").to_pandas(),33        on="arxiv_id",34    )35    df = df[::-1].reset_index(drop=True)36 37    paper_info = []38    for _, row in tqdm.auto.tqdm(df.iterrows(), total=len(df)):39        info = PaperInfo(40            **row,41            paper_page=f"https://huggingface.co/papers/{row.arxiv_id}",42        )43        paper_info.append(info)44    45    df = pd.DataFrame([dataclasses.asdict(info) for info in paper_info])46 47    # set date as index48    df = df.set_index('date')49    df.index = pd.to_datetime(df.index)50    if start_date is not None and end_date is not None:51      # only include data between start_date and end_date52      df = df[(df.index >= start_date) & (df.index <= end_date)]53 54    return df55 56 57def get_github_url(client: PapersWithCodeClient, paper_title: str) -> str:58  """59  Get the Github URL for a paper.60  """61 62  repo_url = ""63  try:64    # get paper ID65    results = client.paper_list(q=paper_title).results66    paper_id = results[0].id67 68    # get paper69    paper = client.paper_get(paper_id=paper_id)70 71    # get repositories72    repositories = client.paper_repository_list(paper_id=paper.id).results73 74    for repo in repositories:75      if repo.is_official:76        repo_url = repo.url77 78  except:79    pass80 81  return repo_url82 83 84def add_metadata_batch(batch, client: PapersWithCodeClient):85    """86    Add metadata to a batch of papers.87    """88 89    # get Github URLs for all papers in the batch90    github_urls = []91    for paper_title in batch["title"]:92        github_url = get_github_url(client, paper_title)93        github_urls.append(github_url)94 95    # overwrite the Github links96    batch["github"] = github_urls97 98    return batch99 100 101def add_hf_assets(batch):102    """103    Add Hugging Face assets to a batch of papers.104    """105    num_spaces = []106    num_models = []107    num_datasets = []108    for arxiv_id in batch["arxiv_id"]:109        if arxiv_id != "":110            response = requests.get(f"https://huggingface.co/api/arxiv/{arxiv_id}/repos")111            result = response.json()112            num_spaces_example = len(result["spaces"])113            num_models_example = len(result["models"])114            num_datasets_example = len(result["datasets"])115        else:116            num_spaces_example = 0117            num_models_example = 0118            num_datasets_example = 0119 120        num_spaces.append(num_spaces_example)121        num_models.append(num_models_example)122        num_datasets.append(num_datasets_example)123 124    batch["num_models"] = num_models125    batch["num_datasets"] = num_datasets126    batch["num_spaces"] = num_spaces127 128    return batch129 130 131def check_hf_mention(batch):132    """133    Check if a paper mentions Hugging Face in the README.134    """135 136    hf_mentions = []137    for github_url in batch["github"]:138        hf_mention = 0139        if github_url != "":140            # get README text using Github API141            owner = github_url.split("/")[-2]142            repo = github_url.split("/")[-1]143            branch = "main"144            url = f"https://raw.githubusercontent.com/{owner}/{repo}/{branch}/README.md"145            response = requests.get(url)146            147            if response.status_code != 200:148                # try master branch as second attempt149                branch = "master"150                url = f"https://raw.githubusercontent.com/{owner}/{repo}/{branch}/README.md"151                response = requests.get(url)152            153            if response.status_code == 200:154                # get text155                text = response.text156                if "huggingface" in text.lower() or "hugging face" in text.lower():157                    hf_mention = 1158            159        hf_mentions.append(hf_mention)160 161    # overwrite the Github links162    batch["hf_mention"] = hf_mentions163 164    return batch165 166 167def process_data(start_date: str, end_date: str) -> pd.DataFrame:168    """169    Load the dataset and enrich it with metadata.170    """171    # step 1. load as HF dataset172    df = get_df(start_date, end_date)173    dataset = Dataset.from_pandas(df)174 175    # step 2. enrich using PapersWithCode API176    dataset = dataset.map(add_metadata_batch, batched=True, batch_size=4, num_proc=cpu_count(), fn_kwargs={"client": PapersWithCodeClient()})177 178    # step 3. enrich using Hugging Face API179    dataset = dataset.map(add_hf_assets, batched=True, batch_size=4, num_proc=cpu_count())180 181    # step 4. check if Hugging Face is mentioned in the README182    dataset = dataset.map(check_hf_mention, batched=True, batch_size=4, num_proc=cpu_count())183 184    # return as Pandas dataframe185    # making sure that the date is set as index186    dataframe = dataset.to_pandas()187    dataframe = dataframe.set_index('date')188    dataframe.index = pd.to_datetime(dataframe.index)189 190    return dataframe191 192 193@st.cache_data194def get_data() -> pd.DataFrame:195 196    # step 1: load pre-processed data197    df = load_dataset("nielsr/daily-papers-enriched", split="train").to_pandas()198    df = df.set_index('date')199    df = df.sort_index()200    df.index = pd.to_datetime(df.index)201 202    # step 2: check how much extra data we need to process203    latest_day = df.iloc[-1].name.strftime('%Y-%m-%d')204    today = pd.Timestamp.today().strftime('%Y-%m-%d')205 206    print("Latest day:", latest_day)207    print("Today:", today)208 209    # step 3: process the missing data210    if latest_day < today:211        print(f"Processing data from {latest_day} to {today}")212        new_df = process_data(start_date=latest_day, end_date=today)213 214        print("Original df:", df.head())215        print("New df:", new_df.head())216 217        df = pd.concat([df, new_df])218 219    df = df.sort_index()220 221    return df