pxiaoer/papers

1import dataclasses2import datetime3import operator4import pathlib5 6import pandas as pd7import requests8import tqdm.auto9 10 11@dataclasses.dataclass(frozen=True)12class PaperInfo:13    date: str14    arxiv_id: str15    github: str16    title: str17    paper_page: str18    upvotes: int19    published_at: str20 21    def __post_init__(self):22        object.__setattr__(self, "published_at", PaperInfo.convert_timestamp(self.published_at))23 24    @staticmethod25    def convert_timestamp(timestamp: str) -> str:26        try:27            return datetime.datetime.strptime(timestamp, "%Y-%m-%dT%H:%M:%S.%fZ").strftime("%Y/%m/%d %H:%M:%S")28        except ValueError:29            return timestamp30 31 32def get_df(path: pathlib.Path | str) -> pd.DataFrame:33    df = pd.read_csv(path, dtype=str).fillna("")34    paper_info = []35    for _, row in tqdm.auto.tqdm(df.iterrows(), total=len(df)):36        res = requests.get(f"https://huggingface.co/api/papers/{row.arxiv_id}").json()37        info = PaperInfo(38            **row,39            title=res["title"],40            paper_page=f"https://huggingface.co/papers/{row.arxiv_id}",41            upvotes=res["upvotes"],42            published_at=res["publishedAt"],43        )44        paper_info.append(info)45    return pd.DataFrame([dataclasses.asdict(info) for info in paper_info])46 47 48class Prettifier:49    @staticmethod50    def get_github_link(link: str) -> str:51        if not link:52            return ""53        return Prettifier.create_link("github", link)54 55    @staticmethod56    def create_link(text: str, url: str) -> str:57        return f'<a href="{url}" target="_blank">{text}</a>'58 59    @staticmethod60    def to_div(text: str | None, category_name: str) -> str:61        if text is None:62            text = ""63        class_name = f"{category_name}-{text.lower()}"64        return f'<div class="{class_name}">{text}</div>'65 66    def __call__(self, df: pd.DataFrame) -> pd.DataFrame:67        df = df.sort_values("arxiv_id", ascending=False).reset_index(drop=True)68        new_rows = []69        for _, row in df.iterrows():70            new_row = dict(row) | {71                "date": Prettifier.create_link(row.date, f"https://huggingface.co/papers?date={row.date}"),72                "paper_page": f'<a href="%s" target="_blank">%s</a>' % ("https://arxiv.org/abs/"+row.arxiv_id,row.arxiv_id),73                "github": self.get_github_link(row.github),74            }75            new_rows.append(new_row)76        return pd.DataFrame(new_rows, columns=df.columns)77 78 79class PaperList:80    COLUMN_INFO = [81        ["date", "markdown"],82        ["paper_page", "markdown"],83        ["title", "str"],84        ["github", "markdown"],85        ["upvotes", "number"],86    ]87 88    def __init__(self, df: pd.DataFrame):89        self.df_raw = df90        self._prettifier = Prettifier()91        self.df_prettified = self._prettifier(df).loc[:, self.column_names]92 93    @property94    def column_names(self):95        return list(map(operator.itemgetter(0), self.COLUMN_INFO))96 97    @property98    def column_datatype(self):99        return list(map(operator.itemgetter(1), self.COLUMN_INFO))100