Add API for loading data from source configurations

lilac/__init__.py

@@ -1,3 +1,27 @@
+from .data.dataset_duckdb import DatasetDuckDB
+from .data.sources.csv_source import CSVDataset
+from .data.sources.default_sources import register_default_sources
+from .data.sources.gmail_source import Gmail
+from .data.sources.huggingface_source import HuggingFaceDataset
+from .data.sources.json_source import JSONDataset
+from .data_loader import create_dataset
+from .db_manager import get_dataset, set_default_dataset_cls
 from .server import start_server, stop_server
+from .signals.default_signals import register_default_signals
 
-__all__ = ['start_server', 'stop_server']
+register_default_sources()
+register_default_signals()
+set_default_dataset_cls(DatasetDuckDB)
+
+__all__ = [
+  'start_server',
+  'stop_server',
+  'create_dataset',
+  'get_dataset',
+
+  # Source configurations.
+  'HuggingFaceDataset',
+  'CSVDataset',
+  'JSONDataset',
+  'Gmail',
+]

lilac/data_loader.py

@@ -18,10 +18,13 @@
 import pandas as pd
 from distributed import Client
 
+from .config import data_path
+from .data.dataset import Dataset
 from .data.dataset_utils import write_items_to_parquet
 from .data.sources.default_sources import register_default_sources
 from .data.sources.source import Source
 from .data.sources.source_registry import resolve_source
+from .db_manager import get_dataset
 from .schema import (
   MANIFEST_FILENAME,
   PARQUET_FILENAME_PREFIX,
@@ -37,6 +40,16 @@
 from .utils import get_dataset_output_dir, log, open_file
 
 
+def create_dataset(
+  namespace: str,
+  dataset_name: str,
+  source_config: Source,
+) -> Dataset:
+  """Load a dataset from a given source configuration."""
+  process_source(data_path(), namespace, dataset_name, source_config)
+  return get_dataset(namespace, dataset_name)
+
+
 def process_source(base_dir: Union[str, pathlib.Path],
                    namespace: str,
                    dataset_name: str,

lilac/router_data_loader.py

@@ -14,16 +14,13 @@
 
 from .auth import get_user_access
 from .config import data_path
-from .data.sources.default_sources import register_default_sources
 from .data.sources.source_registry import get_source_cls, registered_sources
 from .data_loader import process_source
 from .router_utils import RouteErrorHandler
 from .tasks import TaskId, task_manager
 
 REQUEST_TIMEOUT_SEC = 30 * 60  # 30 mins.
 
-register_default_sources()
-
 router = APIRouter(route_class=RouteErrorHandler)
 
 

lilac/router_dataset.py

@@ -23,13 +23,11 @@
   StatsResult,
   UnaryOp,
 )
-from .data.dataset_duckdb import DatasetDuckDB
-from .db_manager import get_dataset, remove_dataset_from_cache, set_default_dataset_cls
+from .db_manager import get_dataset, remove_dataset_from_cache
 from .router_utils import RouteErrorHandler
 from .schema import Bin, Path, normalize_path
 from .signals.concept_labels import ConceptLabelsSignal
 from .signals.concept_scorer import ConceptScoreSignal
-from .signals.default_signals import register_default_signals
 from .signals.semantic_similarity import SemanticSimilaritySignal
 from .signals.signal import (
   Signal,
@@ -44,9 +42,6 @@
 
 router = APIRouter(route_class=RouteErrorHandler)
 
-register_default_signals()
-set_default_dataset_cls(DatasetDuckDB)
-
 
 @router.get('/', response_model_exclude_none=True)
 def get_datasets() -> list[DatasetInfo]:

lilac/server.py

@@ -163,7 +163,7 @@ def start_server(host: str = '0.0.0.0', port: int = 5432) -> None:
   if SERVER:
     raise ValueError('Server is already running')
 
-  config = uvicorn.Config(app, host='0.0.0.0', port=5432)
+  config = uvicorn.Config(app, host='0.0.0.0', port=5432, access_log=False)
   SERVER = uvicorn.Server(config)
   try:
     loop = asyncio.get_running_loop()

notebooks/API.ipynb

@@ -4,111 +4,190 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Starting the web server\n"
+    "## Creating a Lilac dataset\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import lilac as ll"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### From HuggingFace\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/Users/dsmilkov/code/lilac/.venv/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
-      "  from .autonotebook import tqdm as notebook_tqdm\n"
+      "Found cached dataset glue (/Users/dsmilkov/.cache/huggingface/datasets/glue/ax/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad)\n",
+      "100%|██████████| 1/1 [00:00<00:00, 399.38it/s]\n"
      ]
     },
     {
-     "name": "stderr",
+     "name": "stdout",
      "output_type": "stream",
      "text": [
-      "INFO:     Started server process [72797]\n",
-      "INFO:     Waiting for application startup.\n",
-      "INFO:     Application startup complete.\n",
-      "INFO:     Uvicorn running on http://0.0.0.0:5432 (Press CTRL+C to quit)\n"
+      "Manifest for dataset \"glue\" written to ./data/datasets/local/glue\n"
      ]
     }
    ],
    "source": [
-    "import lilac as ll\n",
-    "\n",
-    "ll.start_server()"
+    "source_config = ll.HuggingFaceDataset(dataset_name='glue', config_name='ax')\n",
+    "dataset = ll.create_dataset('local', 'glue', source_config)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Stopping the server\n"
+    "### From CSV\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
-     "name": "stderr",
+     "name": "stdout",
      "output_type": "stream",
      "text": [
-      "INFO:     Shutting down\n",
-      "INFO:     Waiting for application shutdown.\n",
-      "INFO:     Application shutdown complete.\n"
+      "Downloading from url https://storage.googleapis.com/lilac-data-us-east1/datasets/csv_datasets/the_movies_dataset/the_movies_dataset.csv to /tmp/./data/local_cache/37a9be0240c140da95c664d3bc092a04\n",
+      "Manifest for dataset \"the_movies_dataset\" written to ./data/datasets/local/the_movies_dataset\n"
      ]
     }
    ],
    "source": [
-    "await ll.stop_server()"
+    "source_config = ll.CSVDataset(filepaths=[\n",
+    "  'https://storage.googleapis.com/lilac-data-us-east1/datasets/csv_datasets/the_movies_dataset/the_movies_dataset.csv'\n",
+    "])\n",
+    "dataset = ll.create_dataset('local', 'the_movies_dataset', source_config)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Importing a dataset\n"
+    "### From JSON\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Downloading from url https://raw.githubusercontent.com/explosion/prodigy-recipes/master/example-datasets/news_headlines.jsonl to /tmp/./data/local_cache/7f919bed71f0491cae57f3b9cd4f0aa4\n",
+      "Manifest for dataset \"news_headlines\" written to ./data/datasets/local/news_headlines\n"
+     ]
+    }
+   ],
    "source": [
-    "dataset = lilac.load(namespace='local', name=..., config: SourceConfig) # Blocking."
+    "source_config = ll.JSONDataset(filepaths=[\n",
+    "  'https://raw.githubusercontent.com/explosion/prodigy-recipes/master/example-datasets/news_headlines.jsonl'\n",
+    "])\n",
+    "dataset = ll.create_dataset('local', 'news_headlines', source_config)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Query a dataset\n"
+    "## Get an existing dataset\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
-    "dataset.select_rows()  # ......\n"
+    "dataset = ll.get_dataset('local', 'the_movies_dataset')"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Index the dataset\n"
+    "## Visualize the data\n",
+    "\n",
+    "Now that we have imported a few datasets, let's visualize them to see what they look like.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:     Started server process [78275]\n",
+      "INFO:     Waiting for application startup.\n",
+      "INFO:     Application startup complete.\n",
+      "INFO:     Uvicorn running on http://0.0.0.0:5432 (Press CTRL+C to quit)\n"
+     ]
+    }
+   ],
+   "source": [
+    "ll.start_server()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Stopping the server\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:     Shutting down\n",
+      "INFO:     Waiting for application shutdown.\n",
+      "INFO:     Application shutdown complete.\n"
+     ]
+    }
+   ],
+   "source": [
+    "await ll.stop_server()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Query a dataset\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
    "source": [
-    "dataset.compute_embedding(...)"
+    "## Index the dataset\n"
    ]
   },
   {
@@ -147,13 +226,7 @@
     }
    ],
    "source": [
-    "from lilac.db_manager import get_dataset, set_default_dataset_cls\n",
-    "from lilac.data.dataset_duckdb import DatasetDuckDB\n",
     "from lilac.signals.concept_scorer import ConceptScoreSignal\n",
-    "from lilac.signals.default_signals import register_default_signals\n",
-    "\n",
-    "register_default_signals()\n",
-    "set_default_dataset_cls(DatasetDuckDB)\n",
     "\n",
     "dataset = get_dataset('local', 'legal-clauses')\n",
     "\n",
@@ -201,13 +274,6 @@
     "3. I want to download it\n"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
   {
    "cell_type": "markdown",
    "metadata": {},