Switch API to load Hugging Face dataset

app.py

@@ -1,9 +1,8 @@
-import json
 import logging
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
-import pandas as pd
+from datasets import load_dataset
 import numpy as np
 import os
 
@@ -19,7 +18,7 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 # ---------------------------------------------------------
-# CORS (optional but useful for frontend)
+# CORS
 # ---------------------------------------------------------
 app.add_middleware(
     CORSMiddleware,
@@ -29,28 +28,26 @@ app.add_middleware(
 )
 
 # ---------------------------------------------------------
-# Lazy-loaded dataset
+# Load Hugging Face dataset (cached)
 # ---------------------------------------------------------
-DATASET_PATH = "transcripts.csv"   # change if needed
-df_cache = None
+DATASET_NAME = "kurry/sp500_earnings_transcripts"
+dataset_cache = None
 
-def load_dataset():
-    global df_cache
-    if df_cache is None:
-        logger.info("Loading dataset...")
-        if not os.path.exists(DATASET_PATH):
-            raise FileNotFoundError(f"Dataset not found: {DATASET_PATH}")
-        df_cache = pd.read_csv(DATASET_PATH)
-        logger.info(f"Dataset loaded with {len(df_cache)} rows.")
-    return df_cache
+def load_hf_dataset():
+    global dataset_cache
+    if dataset_cache is None:
+        logger.info(f"Loading HF dataset: {DATASET_NAME}")
+        dataset_cache = load_dataset(DATASET_NAME, split="train")
+        logger.info(f"Loaded {len(dataset_cache)} rows from HF dataset")
+    return dataset_cache
 
 # ---------------------------------------------------------
 # JSON-safe conversion
 # ---------------------------------------------------------
 def to_json_safe(obj):
-    if isinstance(obj, (np.integer, np.int64)):
+    if isinstance(obj, (np.integer,)):
         return int(obj)
-    if isinstance(obj, (np.floating, np.float64)):
+    if isinstance(obj, (np.floating,)):
         return float(obj)
     if isinstance(obj, (np.ndarray, list)):
         return [to_json_safe(x) for x in obj]
@@ -59,7 +56,7 @@ def to_json_safe(obj):
     return obj
 
 # ---------------------------------------------------------
-# Serve index.html at "/"
+# Serve index.html
 # ---------------------------------------------------------
 @app.get("/", response_class=HTMLResponse)
 def serve_index():
@@ -73,8 +70,8 @@ def serve_index():
 # ---------------------------------------------------------
 @app.get("/tickers")
 def get_tickers():
-    df = load_dataset()
-    tickers = sorted(df["symbol"].unique().tolist())
+    ds = load_hf_dataset()
+    tickers = sorted(set(ds["symbol"]))
     return {"tickers": tickers}
 
 # ---------------------------------------------------------
@@ -82,19 +79,17 @@ def get_tickers():
 # ---------------------------------------------------------
 @app.get("/transcript/{symbol}")
 def get_transcript(symbol: str):
-    df = load_dataset()
-
+    ds = load_hf_dataset()
     symbol = symbol.upper()
+
     logger.info(f"Fetching transcript for: {symbol}")
 
-    subset = df[df["symbol"] == symbol]
+    rows = [r for r in ds if r["symbol"].upper() == symbol]
 
-    if subset.empty:
+    if not rows:
         raise HTTPException(status_code=404, detail=f"No transcript found for {symbol}")
 
-    # Convert each row to JSON-safe dict
-    records = subset.to_dict(orient="records")
-    safe_records = [to_json_safe(r) for r in records]
+    safe_rows = [to_json_safe(r) for r in rows]
 
-    return {"symbol": symbol, "records": safe_records}
+    return {"symbol": symbol, "records": safe_rows}