Jenny0613/Datasets-in-machine-learning
0
1"""2FastAPI + HTMX app for browsing arxiv papers with new ML datasets.3Downloads Lance dataset from HuggingFace Hub and loads locally.4"""5 6import math7import re8from datetime import date, timedelta9from functools import lru_cache10from typing import Optional11from urllib.parse import urlencode12 13import lance14import polars as pl15from cachetools import TTLCache16from dotenv import load_dotenv17from fastapi import FastAPI, Query, Request18from fastapi.responses import HTMLResponse, RedirectResponse19from fastapi.staticfiles import StaticFiles20from fastapi.templating import Jinja2Templates21from huggingface_hub import snapshot_download22from markupsafe import Markup23 24# Load .env file for local development (HF_TOKEN)25load_dotenv()26 27app = FastAPI(title="ArXiv New ML Datasets")28app.mount("/static", StaticFiles(directory="static"), name="static")29templates = Jinja2Templates(directory="templates")30 31 32def highlight_search(text: str, search: str) -> Markup:33 """Highlight search terms in text with yellow background."""34 if not search or not text:35 return Markup(text) if text else Markup("")36 37 # Escape HTML in text first38 import html39 text = html.escape(str(text))40 41 # Case-insensitive replacement with highlight span42 pattern = re.compile(re.escape(search), re.IGNORECASE)43 highlighted = pattern.sub(44 lambda m: f'<mark class="bg-yellow-200 px-0.5 rounded">{m.group()}</mark>',45 text46 )47 return Markup(highlighted)48 49 50# Register custom filters51templates.env.filters["highlight"] = highlight_search52 53 54def confidence_fmt(score):55 """Format confidence as percentage, truncating to 1 decimal to avoid rounding 99.95->100."""56 pct = math.floor(score * 1000) / 1057 return f"{pct:.1f}"58 59 60templates.env.filters["confidence"] = confidence_fmt61 62# Dataset config63DATASET_REPO = "librarian-bots/arxiv-cs-papers-lance"64 65# Cache for dataset (reload every 6 hours)66_dataset_cache: TTLCache = TTLCache(maxsize=1, ttl=60 * 60 * 6)67 68# Cache for Lance dataset connection (for vector search)69_lance_cache: dict = {}70 71# Cache for embedding model (lazy loaded on first semantic search)72_model_cache: dict = {}73 74 75def get_lance_dataset():76 """Download dataset from HF Hub (cached) and return Lance connection."""77 if "ds" not in _lance_cache:78 import os79 # Use HF_HOME or /tmp for Spaces compatibility (./data not writable on Spaces)80 cache_base = os.environ.get("HF_HOME", "/tmp/hf_cache")81 local_dir = f"{cache_base}/arxiv-lance"82 print(f"Downloading dataset from {DATASET_REPO} to {local_dir}...")83 snapshot_download(84 DATASET_REPO,85 repo_type="dataset",86 local_dir=local_dir,87 )88 lance_path = f"{local_dir}/data/train.lance"89 print(f"Loading Lance dataset from {lance_path}")90 _lance_cache["ds"] = lance.dataset(lance_path)91 return _lance_cache["ds"]92 93 94def get_embedding_model():95 """Load embedding model (cached, lazy-loaded on first semantic search)."""96 if "model" not in _model_cache:97 from sentence_transformers import SentenceTransformer98 print("Loading embedding model...")99 _model_cache["model"] = SentenceTransformer("BAAI/bge-base-en-v1.5")100 print("Embedding model loaded!")101 return _model_cache["model"]102 103 104def get_dataframe() -> pl.DataFrame:105 """Load Lance dataset and convert to Polars DataFrame."""106 cache_key = "df"107 if cache_key in _dataset_cache:108 return _dataset_cache[cache_key]109 110 ds = get_lance_dataset() # Downloads from HF Hub if not cached111 # Select columns needed for filtering/display (exclude embeddings for memory)112 columns = [113 "id", "title", "abstract", "categories", "update_date",114 "authors", "is_new_dataset", "confidence_score"115 ]116 arrow_table = ds.to_table(columns=columns)117 df = pl.from_arrow(arrow_table)118 _dataset_cache[cache_key] = df119 print(f"Loaded {len(df):,} papers")120 return df121 122 123@lru_cache(maxsize=1)124def get_categories() -> list[str]:125 """Get unique category prefixes for filtering."""126 df = get_dataframe()127 # Extract primary category (before first space or as-is)128 categories = (129 df.select(pl.col("categories").str.split(" ").list.first().alias("cat"))130 .unique()131 .sort("cat")132 .to_series()133 .to_list()134 )135 # Get common ML-related categories136 ml_cats = ["cs.AI", "cs.CL", "cs.CV", "cs.LG", "cs.NE", "cs.IR", "cs.RO", "stat.ML"]137 return [c for c in ml_cats if c in categories]138 139 140@lru_cache(maxsize=1)141def get_confidence_options() -> list[dict]:142 """Compute confidence filter options from actual data distribution.143 144 Uses percentiles so the UI adapts to any model's score range.145 """146 df = get_dataframe()147 scores = df.filter(pl.col("is_new_dataset"))["confidence_score"]148 149 options = [{"value": "0.5", "label": "All new datasets", "count": len(scores)}]150 151 for pct_label, quantile in [("Top 75%", 0.25), ("Top 50%", 0.50), ("Top 25%", 0.75)]:152 threshold = float(scores.quantile(quantile))153 count = scores.filter(scores >= threshold).len()154 options.append({155 "value": f"{threshold:.2f}",156 "label": pct_label,157 "count": int(count),158 })159 160 options.append({"value": "0", "label": "All papers", "count": len(df)})161 return options162 163 164@lru_cache(maxsize=1)165def get_histogram_data() -> dict:166 """Get confidence distribution data for histogram display.167 168 Dynamically determines the range from actual data distribution.169 Returns dict with bins and metadata. The 50% line marks the prediction boundary.170 """171 df = get_dataframe()172 173 # Get all papers with confidence scores174 all_papers = df.select("confidence_score", "is_new_dataset")175 176 # Dynamically determine the range from actual data177 # Round to nearest 5% for clean boundaries178 actual_min = float(all_papers["confidence_score"].min())179 actual_max = float(all_papers["confidence_score"].max())180 181 # Round down to nearest 5% for min, round up for max182 min_pct = max(0, (int(actual_min * 20) / 20)) # Floor to 5%183 max_pct = min(1, ((int(actual_max * 20) + 1) / 20)) # Ceil to 5%184 185 # Ensure minimum range of 25% for usability186 if max_pct - min_pct < 0.25:187 center = (min_pct + max_pct) / 2188 min_pct = max(0, center - 0.125)189 max_pct = min(1, center + 0.125)190 191 # Use 25 bins for good granularity192 num_bins = 25193 bin_width = (max_pct - min_pct) / num_bins194 195 bins = []196 for i in range(num_bins):197 bin_start = min_pct + i * bin_width198 bin_end = min_pct + (i + 1) * bin_width199 200 # Count papers in this bin201 count = all_papers.filter(202 (pl.col("confidence_score") >= bin_start) &203 (pl.col("confidence_score") < bin_end)204 ).height205 206 # Count new_dataset papers in this bin207 new_dataset_count = all_papers.filter(208 (pl.col("confidence_score") >= bin_start) &209 (pl.col("confidence_score") < bin_end) &210 (pl.col("is_new_dataset"))211 ).height212 213 bins.append({214 "bin_start": round(bin_start, 3),215 "bin_end": round(bin_end, 3),216 "bin_pct": int(bin_start * 100),217 "count": count,218 "new_dataset_count": new_dataset_count,219 })220 221 # Normalize counts for display (max height = 100%)222 max_count = max(b["count"] for b in bins) if bins else 1223 for b in bins:224 b["height_pct"] = int((b["count"] / max_count) * 100) if max_count > 0 else 0225 b["new_height_pct"] = int((b["new_dataset_count"] / max_count) * 100) if max_count > 0 else 0226 227 # Calculate cumulative counts from each threshold228 # (how many papers are at or above this threshold)229 total_so_far = all_papers.height230 for b in bins:231 b["papers_above"] = total_so_far232 total_so_far -= b["count"]233 234 return {235 "bins": bins,236 "min_pct": round(min_pct, 2),237 "max_pct": round(max_pct, 2),238 "total_papers": all_papers.height,239 "new_dataset_count": all_papers.filter(pl.col("is_new_dataset")).height,240 }241 242 243def parse_since(since: str) -> Optional[date]:244 """Parse 'since' parameter to a date. Returns None for 'all time'."""245 if not since:246 return None247 today = date.today()248 if since == "1m":249 return today - timedelta(days=30)250 elif since == "6m":251 return today - timedelta(days=180)252 elif since == "1y":253 return today - timedelta(days=365)254 return None255 256 257def filter_papers(258 df: pl.DataFrame,259 category: Optional[str] = None,260 search: Optional[str] = None,261 min_confidence: float = 0.5,262 since: Optional[str] = None,263) -> pl.DataFrame:264 """Apply filters to the papers dataframe.265 266 The confidence threshold controls which papers are shown:267 - Papers with is_new_dataset=True have confidence >= 0.5268 - Setting threshold to 0 shows all papers269 - Setting threshold >= 0.5 effectively shows only new_dataset papers270 """271 if min_confidence >= 0.5:272 # Show only papers classified as new datasets, filtered by confidence273 df = df.filter(274 pl.col("is_new_dataset") & (pl.col("confidence_score") >= min_confidence)275 )276 elif min_confidence > 0:277 df = df.filter(pl.col("confidence_score") >= min_confidence)278 279 if category:280 df = df.filter(pl.col("categories").str.contains(category))281 282 if search:283 search_lower = search.lower()284 df = df.filter(285 pl.col("title").str.to_lowercase().str.contains(search_lower)286 | pl.col("abstract").str.to_lowercase().str.contains(search_lower)287 )288 289 # Date filter290 min_date = parse_since(since)291 if min_date:292 df = df.filter(pl.col("update_date") >= min_date)293 294 return df295 296 297def paginate_papers(298 df: pl.DataFrame,299 page: int = 1,300 per_page: int = 20,301 sort: str = "date",302) -> tuple[pl.DataFrame, bool]:303 """Sort and paginate papers, return (page_df, has_more).304 305 Sort options:306 - "date": By update_date desc, then confidence_score desc307 - "relevance": Keep existing order (for semantic search similarity)308 """309 if sort == "date":310 df_sorted = df.sort(311 ["update_date", "confidence_score"], descending=[True, True]312 )313 else:314 # "relevance" - keep existing order (already sorted by similarity for semantic)315 df_sorted = df316 317 start = (page - 1) * per_page318 page_df = df_sorted.slice(start, per_page + 1)319 has_more = len(page_df) > per_page320 321 return page_df.head(per_page), has_more322 323 324def semantic_search(325 query: str,326 k: int = 100,327 category: Optional[str] = None,328 min_confidence: float = 0.5,329 since: Optional[str] = None,330) -> pl.DataFrame:331 """Search using vector similarity via Lance nearest neighbor.332 333 Returns DataFrame with similarity_score column (0-1, higher is more similar).334 """335 model = get_embedding_model()336 query_embedding = model.encode(query).tolist()337 338 ds = get_lance_dataset()339 340 # Build SQL filter (Lance supports SQL-like syntax)341 filters = []342 if min_confidence >= 0.5:343 filters.append("is_new_dataset = true")344 filters.append(f"confidence_score >= {min_confidence}")345 elif min_confidence > 0:346 filters.append(f"confidence_score >= {min_confidence}")347 if category:348 # Escape single quotes in category name for SQL safety349 safe_category = category.replace("'", "''")350 filters.append(f"categories LIKE '%{safe_category}%'")351 # Date filter - use TIMESTAMP literal for Lance/DataFusion352 min_date = parse_since(since)353 if min_date:354 filters.append(f"update_date >= TIMESTAMP '{min_date.isoformat()} 00:00:00'")355 filter_str = " AND ".join(filters) if filters else None356 357 # Vector search - include _distance for similarity calculation358 results = ds.scanner(359 nearest={"column": "embedding", "q": query_embedding, "k": k},360 filter=filter_str,361 columns=["id", "title", "abstract", "categories", "update_date",362 "authors", "confidence_score", "_distance"]363 ).to_table()364 365 df = pl.from_arrow(results)366 367 # Convert L2 distance to similarity score (0-1 range)368 # For normalized embeddings: similarity = 1 - distance/2369 # BGE embeddings are normalized, so L2 distance ranges from 0 to 2370 df = df.with_columns(371 (1 - pl.col("_distance") / 2).clip(0, 1).alias("similarity_score")372 ).drop("_distance")373 374 return df375 376 377@app.get("/", response_class=HTMLResponse)378async def home(379 request: Request,380 search: Optional[str] = Query(None),381 search_type: str = Query("keyword"),382 category: Optional[str] = Query(None),383 min_confidence: str = Query("0.5"), # String to preserve exact value for template384 since: Optional[str] = Query(None),385 sort: str = Query("date"),386):387 """Render the home page with optional initial filter state from URL."""388 df = get_dataframe()389 categories = get_categories()390 histogram_data = get_histogram_data()391 confidence_options = get_confidence_options()392 393 # Get stats394 total_papers = len(df)395 new_dataset_count = df.filter(pl.col("is_new_dataset")).height396 397 return templates.TemplateResponse(398 "index.html",399 {400 "request": request,401 "categories": categories,402 "total_papers": total_papers,403 "new_dataset_count": new_dataset_count,404 "histogram_data": histogram_data,405 "confidence_options": confidence_options,406 # Pass filter state for URL persistence407 "search": search or "",408 "search_type": search_type,409 "category": category or "",410 "min_confidence": min_confidence,411 "since": since or "",412 "sort": sort,413 },414 )415 416 417@app.get("/papers", response_class=HTMLResponse)418async def get_papers(419 request: Request,420 page: int = Query(1, ge=1),421 per_page: int = Query(20, ge=1, le=100),422 category: Optional[str] = Query(None),423 search: Optional[str] = Query(None),424 min_confidence: float = Query(0.5, ge=0, le=1),425 search_type: str = Query("keyword"), # "keyword" or "semantic"426 sort: str = Query("date"), # "date" or "relevance"427 since: Optional[str] = Query(None), # "1m", "6m", "1y", or None for all428):429 """Get paginated and filtered papers (returns HTML partial for HTMX).430 431 If accessed directly (not via HTMX), redirects to home page with same params.432 """433 # Redirect direct browser visits to home page (this endpoint returns partials)434 if "HX-Request" not in request.headers:435 # Build redirect URL with current query params436 query_string = str(request.url.query)437 redirect_url = f"/?{query_string}" if query_string else "/"438 return RedirectResponse(url=redirect_url, status_code=302)439 440 if search and search_type == "semantic":441 # Vector search - returns pre-sorted by similarity442 filtered_df = semantic_search(443 query=search,444 k=per_page * 5, # Get more for pagination buffer445 category=category,446 min_confidence=min_confidence,447 since=since,448 )449 # Default to relevance sort for semantic, but allow date sort450 effective_sort = sort if sort == "date" else "relevance"451 page_df, has_more = paginate_papers(452 filtered_df, page=page, per_page=per_page, sort=effective_sort453 )454 else:455 # Existing keyword search path456 df = get_dataframe()457 filtered_df = filter_papers(458 df,459 category=category,460 search=search,461 min_confidence=min_confidence,462 since=since,463 )464 # Keyword search always sorts by date465 page_df, has_more = paginate_papers(466 filtered_df, page=page, per_page=per_page, sort="date"467 )468 469 # Convert to list of dicts for template470 papers = page_df.to_dicts()471 472 # Build clean URL for browser history (/ instead of /papers)473 # Only include non-default values to keep URLs short474 params = {}475 if search:476 params["search"] = search477 if search_type != "keyword":478 params["search_type"] = search_type479 if category:480 params["category"] = category481 if min_confidence != 0.5:482 params["min_confidence"] = min_confidence483 if since:484 params["since"] = since485 if sort != "date":486 params["sort"] = sort487 push_url = "/?" + urlencode(params) if params else "/"488 489 response = templates.TemplateResponse(490 "partials/paper_list.html",491 {492 "request": request,493 "papers": papers,494 "page": page,495 "has_more": has_more,496 "category": category or "",497 "search": search or "",498 "min_confidence": min_confidence,499 "search_type": search_type,500 "sort": sort,501 "since": since or "",502 "total_filtered": len(filtered_df),503 },504 )505 # Tell HTMX to push clean URL (/ not /papers)506 response.headers["HX-Push-Url"] = push_url507 return response508 509 510@app.get("/api/stats")511async def get_stats():512 """Get dataset statistics as JSON."""513 df = get_dataframe()514 515 new_datasets = df.filter(pl.col("is_new_dataset"))516 517 return {518 "total_papers": len(df),519 "new_dataset_count": len(new_datasets),520 "avg_confidence": float(df["confidence_score"].mean()),521 "date_range": {522 "min": str(df["update_date"].min()),523 "max": str(df["update_date"].max()),524 },525 }526 527 528# Preload dataset and model on startup529@app.on_event("startup")530async def startup_event():531 """Preload dataset and embedding model on startup."""532 print("Preloading dataset...")533 get_dataframe()534 print("Dataset loaded!")535 print("Preloading embedding model...")536 get_embedding_model()537 print("Embedding model loaded!")538 539 540if __name__ == "__main__":541 import uvicorn542 543 uvicorn.run(app, host="0.0.0.0", port=7860)544