Asif Rahman

ChartQL - A token-efficient SQL extension for generating Plotly charts

Fri, 12 Sep 2025 00:00:00 +0000

Most LLM-generated chart code is wasteful (“token inefficient”). When you ask a model to create a Plotly chart, it writes mostly boilerplate with a lot of output tokens for what is fundamentally a simple mapping: take these columns, plot them as this chart type, apply these visual options.

ChartQL compresses that entire specification into a few lines. The model only needs to produce a SQL query with a PLOT AS clause, and the execution engine handles the rest. A experiments I found typical Plotly chart specification runs 400-600 tokens. The equivalent ChartQL is 40-80, which is a 10x reduction.

This is especially important in agentic systems where we want to reduce the number of tool calls. The standard flow would be to write SQL to get the data, then write the chart specification. With ChartQL, a single tool call handles both the query and the chart definition. Fewer tool calls means fewer round trips, lower latency, and less opportunity for the model to introduce errors between steps.

Plotly has an excellent reference doc and a machine readable specification that describes all the options for layouts and formatting a chart. But with so many options, a full Plotly specification has many ways it can go wrong, mismatched list lengths, incorrect nesting of layout dicts, and wrong trace types. ChartQL’s grammar is small enough that a model can reliably produce valid output, and when it doesn’t, the parser gives precise error locations in the chart definition, so it’s easy to fix.

ChartQL allows you to embed Plotly chart definitions directly within SQL queries using the PLOT AS syntax.

You can customize chart appearance using Plotly layout options using Plotly’s magic underscore notation which makes it easier to work with nested properties. Supported plot types:

line - Line charts with markers
bar - Bar charts (vertical/horizontal)
scatter - Scatter plots
pie - Pie charts
heatmap - Heat maps
histogram - Histogram plots
box - Box plots

SELECT end_date, n_reports FROM metrics
WHERE measure = 'active_users' AND adjustment = "total"
PLOT AS line(
 x=end_date,
 y=n_reports,
 trace_marker_color='blue',
 trace_marker_size=8,
 layout_title_text='Active Users Over Time',
 layout_xaxis_type='date',
 layout_xaxis_title='Date',
 layout_yaxis_title='Number of Active Users',
 layout_height=400,
 layout_width=700,
 layout_margin_b=50
)

Overall architecture:

sqlparse tokenizes SQL query to identify PLOT AS clauses.
The SQL query is executed to fetch data. The parser does not execute the query itself but relies on an execution engine.
The chart definition is parsed to identify the chart type and its parameters.
Chart generator returns a Plotly chart specification based on the data and chart type: {data, layout, config}.
- Anything with layout_ or config_ prefix is treated as a layout or config option
- Column names can be used directly as variables in the chart definition, like line(x=column_name, y=other_column_name)
  - Only in fields which are expected to be data fields, including: x, y, z, color, values, labels, text, etc.
  - Multiple columns can be used, e.g. y=[col1, col2, col3] for multiple traces
- Trace specific options can be set using the trace_ prefix, e.g. trace_marker_color, trace_marker_size, etc. will add the options to each traces marker dict.
  - When multiple columns are used, trace options can be specified as lists to apply to each trace, like trace_marker_color=[‘red’, ‘blue’, ‘green’]

import sqlparse
import polars as pl
from sqlparse.tokens import Text
from lark import Lark, Transformer
from lark.exceptions import UnexpectedToken
from typing import Dict, Any, Tuple, List, Union, Callable

class ChartQLException(Exception):
 """Base exception for ChartQL parsing and processing errors."""
 pass

class ChartQLParseError(ChartQLException):
 """Enhanced parsing error with visual context showing the error location."""

 def __init__(self, original_error: UnexpectedToken, chart_spec: str):
 self.original_error = original_error
 self.chart_spec = chart_spec
 super().__init__(self._format_error())

 def _format_error(self) -> str:
 lines = self.chart_spec.splitlines()
 error_line_idx = self.original_error.line - 1 # Convert to 0-based
 error_col = self.original_error.column - 1

 # Get the original error message
 message = str(self.original_error)

 # Add visual context if we can locate the line
 if 0 <= error_line_idx < len(lines):
 line_content = lines[error_line_idx]
 pointer = " " * error_col + "^"
 message += f"\n{line_content}\n{pointer}"

 return message


def find_plot_as_position(tokens: List) -> int:
 """Find the position of PLOT AS clause in token list. Returns -1 if not found."""
 for i in range(len(tokens) - 2):
 if (
 tokens[i].value.upper() == "PLOT"
 and tokens[i + 1].ttype in (None, Text.Whitespace)
 and tokens[i + 2].value.upper() == "AS"
 ):
 return i
 return -1


def split_chartql_query(query: str) -> Tuple[str, str]:
 """Split ChartQL query into SQL and chart specification parts."""
 tokens = list(sqlparse.parse(query)[0].flatten())
 plot_as_pos = find_plot_as_position(tokens)

 if plot_as_pos == -1:
 raise ChartQLException("No PLOT AS clause found")

 sql_tokens = tokens[:plot_as_pos]
 chart_tokens = tokens[plot_as_pos + 3 :]

 sql_query = "".join(token.value for token in sql_tokens).strip()
 chart_spec = "".join(token.value for token in chart_tokens).strip()

 return sql_query, chart_spec


def has_chartql_plot(query: str) -> bool:
 """Check if query contains a PLOT AS clause."""
 tokens = list(sqlparse.parse(query)[0].flatten())
 return find_plot_as_position(tokens) != -1


chartql_grammar = """
 start: chart_type "(" parameter_list? ")"
 chart_type: CHART_TYPE
 parameter_list: parameter ("," parameter)*
 parameter: key "=" value
 key: CNAME
 value: string | number | boolean | list | column_ref
 list: "[" value ("," value)* "]"
 string: SINGLE_STRING | DOUBLE_STRING
 number: SIGNED_NUMBER
 boolean: "True" | "False" | "true" | "false"
 column_ref: CNAME

 CHART_TYPE: "line" | "bar" | "scatter" | "pie" | "heatmap" | "histogram" | "box"
 SINGLE_STRING: /'[^']*'/
 DOUBLE_STRING: /"[^"]*"/
 COMMENT: /--[^\\n\\r]*/

 %import common.SIGNED_NUMBER
 %import common.CNAME
 %import common.WS
 %ignore WS
 %ignore COMMENT
"""


class ChartQLTransformer(Transformer):
 def start(self, items):
 chart_type, params = items[0], items[1] if len(items) > 1 else {}
 return {"chart_type": chart_type, "parameters": params}

 def chart_type(self, items):
 return str(items[0])

 def parameter_list(self, items):
 return dict(items)

 def parameter(self, items):
 return (items[0], items[1])

 def key(self, items):
 return str(items[0])

 def list(self, items):
 return list(items)

 def string(self, items):
 return str(items[0])[1:-1]

 def number(self, items):
 val = str(items[0])
 return int(val) if "." not in val else float(val)

 def boolean(self, items):
 if not items:
 return False
 return str(items[0]).lower() == "true"

 def column_ref(self, items):
 return str(items[0])

 def value(self, items):
 return items[0]


def parse_chartql(chart_spec: str) -> Dict[str, Any]:
 """Parse chart specification into structured dictionary."""
 parser = Lark(chartql_grammar, parser="lalr", transformer=ChartQLTransformer())
 try:
 return parser.parse(chart_spec)
 except UnexpectedToken as e:
 raise ChartQLParseError(e, chart_spec) from None


DataInput = Union[List[Dict[str, Any]], pl.DataFrame]


class ChartSpecGenerator:
 def __init__(self):
 self.chart_defaults = {
 "line": {"type": "scatter", "mode": "lines+markers"},
 "bar": {"type": "bar"},
 "scatter": {"type": "scatter", "mode": "markers"},
 "pie": {"type": "pie"},
 "heatmap": {"type": "heatmap"},
 "histogram": {"type": "histogram"},
 "box": {"type": "box"},
 }

 def generate_spec(self, parsed_chart: Dict[str, Any], data: DataInput) -> Dict[str, Any]:
 chart_type = parsed_chart["chart_type"]
 parameters = parsed_chart["parameters"]

 layout_params = self._extract_prefixed_params(parameters, "layout_")
 config_params = self._extract_prefixed_params(parameters, "config_")
 trace_params = self._extract_prefixed_params(parameters, "trace_")
 data_params = self._extract_data_params(parameters)

 traces = self._build_traces(chart_type, data_params, trace_params, data)
 layout = self._build_layout(layout_params)
 config = self._build_config(config_params)

 return {"data": traces, "layout": layout, "config": config}

 def _extract_prefixed_params(self, parameters: Dict[str, Any], prefix: str) -> Dict[str, Any]:
 result = {}
 for key, value in parameters.items():
 if key.startswith(prefix):
 new_key = key[len(prefix) :]
 result[new_key] = value
 return result

 def _extract_data_params(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
 data_fields = {"x", "y", "z", "color", "values", "labels", "text", "size"}
 result = {}
 for key, value in parameters.items():
 if key in data_fields:
 result[key] = value
 return result

 def _build_traces(
 self, chart_type: str, data_params: Dict[str, Any], trace_params: Dict[str, Any], data: DataInput
 ) -> List[Dict[str, Any]]:
 base_trace = self.chart_defaults[chart_type].copy()

 primary_data_col = self._get_primary_data_column(chart_type)
 data_columns = data_params.get(primary_data_col, [])

 if not isinstance(data_columns, list):
 data_columns = [data_columns]

 traces = []
 for i, col in enumerate(data_columns):
 trace = base_trace.copy()

 self._apply_data_mappings(trace, chart_type, data_params, data, i)
 self._apply_trace_params(trace, trace_params, i)

 trace["name"] = col if len(data_columns) > 1 else chart_type
 traces.append(trace)

 return traces

 def _apply_trace_params(self, trace: Dict[str, Any], trace_params: Dict[str, Any], trace_index: int):
 for key, value in trace_params.items():
 if isinstance(value, list) and len(value) > trace_index:
 self._set_nested_dict(trace, key, value[trace_index])
 elif not isinstance(value, list):
 self._set_nested_dict(trace, key, value)

 def _apply_data_mappings(
 self, trace: Dict[str, Any], chart_type: str, data_params: Dict[str, Any], data: DataInput, trace_index: int
 ):
 data_field_mappings = {
 "line": ["x", "y"],
 "scatter": ["x", "y", "size", "color", "text"],
 "bar": ["x", "y"],
 "pie": ["values", "labels"],
 "heatmap": ["x", "y", "z"],
 "histogram": ["x"],
 "box": ["x", "y"],
 }

 for field in data_field_mappings.get(chart_type, []):
 if field in data_params:
 col_ref = data_params[field]
 if isinstance(col_ref, list):
 primary_col = self._get_primary_data_column(chart_type)
 if field == primary_col and len(col_ref) > trace_index:
 trace[field] = self._get_column_data(data, col_ref[trace_index])
 else:
 trace[field] = self._get_column_data(data, col_ref)

 def _get_primary_data_column(self, chart_type: str) -> str:
 primary_cols = {"line": "y", "scatter": "y", "bar": "y", "pie": "values", "histogram": "x", "box": "y"}
 return primary_cols.get(chart_type, "y")

 def _get_column_data(self, data: DataInput, column: str) -> List[Any]:
 """Extract column data as list from either Polars DataFrame or list of dicts."""
 if pl and isinstance(data, pl.DataFrame):
 return data[column].to_list()
 elif isinstance(data, list):
 return [row.get(column) for row in data]
 else:
 raise ChartQLException(f"Unsupported data type: {type(data)}")

 def _get_columns(self, data: DataInput) -> List[str]:
 """Get column names from data."""
 if pl and isinstance(data, pl.DataFrame):
 return data.columns
 elif isinstance(data, list) and data:
 return list(data[0].keys())
 else:
 return []

 def _build_layout(self, layout_params: Dict[str, Any]) -> Dict[str, Any]:
 layout = {}
 for key, value in layout_params.items():
 self._set_nested_dict(layout, key, value)
 return layout

 def _build_config(self, config_params: Dict[str, Any]) -> Dict[str, Any]:
 return config_params.copy()

 def _set_nested_dict(self, d: Dict[str, Any], key: str, value: Any):
 if not key:
 return
 keys = key.split("_")
 if not keys:
 return
 current = d
 for k in keys[:-1]:
 if k not in current:
 current[k] = {}
 current = current[k]
 current[keys[-1]] = value


class DataMapper:
 def __init__(self, data: DataInput):
 self.data = data
 self.columns = set(self._get_columns(data))

 def resolve_column_references(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
 resolved = {}
 for key, value in parameters.items():
 if self._is_data_field(key):
 resolved[key] = self._resolve_value(value)
 else:
 resolved[key] = value
 return resolved

 def _is_data_field(self, key: str) -> bool:
 data_fields = {"x", "y", "z", "color", "values", "labels", "text", "size"}
 return key in data_fields

 def _resolve_value(self, value: Any) -> Any:
 if isinstance(value, str) and value in self.columns:
 return value
 elif isinstance(value, list):
 return [self._resolve_value(v) for v in value]
 else:
 return value

 def validate_columns(self, parameters: Dict[str, Any]):
 missing_columns = []
 for key, value in parameters.items():
 if self._is_data_field(key):
 columns = self._extract_column_names(value)
 for col in columns:
 if col not in self.columns:
 missing_columns.append(col)

 if missing_columns:
 raise ChartQLException(f"Columns not found in data: {missing_columns}")

 def _extract_column_names(self, value: Any) -> List[str]:
 if isinstance(value, str):
 return [value] if value in self.columns else []
 elif isinstance(value, list):
 columns = []
 for v in value:
 columns.extend(self._extract_column_names(v))
 return columns
 else:
 return []

 def _get_columns(self, data: DataInput) -> List[str]:
 """Get column names from data."""
 if pl and isinstance(data, pl.DataFrame):
 return data.columns
 elif isinstance(data, list) and data:
 return list(data[0].keys())
 else:
 return []


def process_chartql_query(query: str, sql_executor: Callable) -> Dict[str, Any]:
 """Process a complete ChartQL query and return Plotly specification."""
 sql_query, chart_spec = split_chartql_query(query)

 data = sql_executor(sql_query)

 parsed_chart = parse_chartql(chart_spec)

 data_mapper = DataMapper(data)
 data_mapper.validate_columns(parsed_chart["parameters"])
 resolved_params = data_mapper.resolve_column_references(parsed_chart["parameters"])

 spec_generator = ChartSpecGenerator()
 plotly_spec = spec_generator.generate_spec(
 {"chart_type": parsed_chart["chart_type"], "parameters": resolved_params}, data
 )

 return plotly_spec

Installing Local Python Packages with uv tool

Mon, 14 Jul 2025 00:00:00 +0000

Astral’s uv Python package manager has a concept of tools, which are Python packages that provide command-line interfaces. Tools are installed with isolated dependencies. For example, I can install my presskit package as a tool and it will be available system-wide: uv tool install presskit.

During development, I often want to test changes to my package without having to publish it to PyPI. uv tool allows me to install a local package or a built distribution directly, making it easy to test changes. Here are a few ways to do it:

Install from local directory

uv tool install .

Install from built distribution

# First build the package
source .venv/bin/activate && python -m build

# Then install the built wheel
uv tool install dist/presskit-*.whl

Install in editable mode for development

uv tool install --editable .

# Install with specific features
uv tool install --editable ".[dev,docs]"

The --editable flag is particularly useful during development as changes to my source code will be reflected immediately without reinstalling.

Install from git repository

uv tool install git+https://github.com/asifr/presskit.git

Disk-based caching with asset management in Python

Tue, 17 Jun 2025 00:00:00 +0000

The Python DiskCache library is my go-to solution for disk-based caching. It uses a SQlite database to store small data, and stores larger data in pickle files on disk. It has expiration and eviction policies and a convenient @memoize decorator for caching function results.

I use Polars and Numpy for data processing and there are more efficient native formats for storing data, such as Parquet, Arrow, and Numpy. Below is an extension to DiskCache that adds support for these formats using a new @asset decorator. Now the Cache class will use duck-typing to determine the type of data being returned by the decorated function and store it in the appropriate format.

cache = Cache(directory="./cache")

@cache.asset()
def filter_weather_stations(df: pl.DataFrame) -> pl.DataFrame:
 return df.filter(pl.col("country_code") == "US")

The @asset decorator will automatically detect the return data is a Polars DataFrame and store it in .parquet format. PyArrow tables are stored as .arrow as Arrow IPC format. Numpy array will be stored in .npy format. If the data is a dictionary or list, it will be stored in JSON format. For any other data type, it will fall back to using Pickle. The Cache class can be extended to support more formats.

This is a complete-reimplementation of the DiskCache core Cache class to add the assets decorator, additional bug-fixes and improvements like hardening against edge cases, including race conditions, file name collisions, and thread safety. The complete implementation is at the end of this note.

Table of Contents:

Features
Basic Usage
Expiration Examples
Memoization Decorator
Asset Handling with Native Formats
Extensible Asset Handlers
Clearing Cache by Name or Function
Storage Backends
Eviction Policies
Pickle Fallback
Advanced Configuration
Implementation code

Features

Disk and SQLite Storage: Hybrid storage using SQLite database for metadata and filesystem for large assets
Expiration Support: Set TTL (time-to-live) for cached items with automatic cleanup
Memoization: Function result caching with the @memoize decorator
Asset Handling: Store data in native formats (JSON, Parquet, Arrow, NumPy, etc.)
Extensible Handlers: Custom asset format handlers with automatic fallback to pickle
Key Management: Clear specific keys, function caches, or all items
Eviction Policies: LRU, LFU, least-recently-stored, or none
Transaction Safety: Thread and process-safe operations with SQLite WAL mode

Basic Usage

Create a cache instance and store/retrieve data:

from explore.cache import Cache

# Create cache with default settings
cache = Cache()

# Store with expiration (3600 seconds = 1 hour)
cache.set("user:123", {"name": "John", "email": "john@example.com"}, expire=3600)

# Retrieve data
user_data = cache.get("user:123")
print(user_data) # {'name': 'John', 'email': 'john@example.com'}

# Check if key exists
if "user:123" in cache:
 print("User data is cached")

# Dictionary-style access
cache["session:abc"] = "session_data"
session = cache["session:abc"]

Expiration Examples

import time

# Set with 5 second expiration
cache.set("temp_key", "temporary_value", expire=5)
print(cache.get("temp_key")) # "temporary_value"

time.sleep(6)
print(cache.get("temp_key")) # None (expired)

# Touch to extend expiration
cache.set("extend_key", "value", expire=10)
cache.touch("extend_key", expire=3600) # Extend to 1 hour

Memoization Decorator

Cache expensive function results automatically:

@cache.memoize(expire=3600, typed=True)
def fibonacci(n):
 if n <= 1:
 return n
 return fibonacci(n-1) + fibonacci(n-2)

# First call computes and caches
result1 = fibonacci(100)

# Subsequent calls return cached result
result2 = fibonacci(100) # Much faster

# Access cache key for manual operations
key = fibonacci.__cache_key__(100)
cached_value = cache[key]

# Clear specific function's cache
cache.clear(fibonacci)

Asset Handling with Native Formats

Store data in native formats instead of pickle:

import pandas as pd
import polars as pl
import numpy as np
import pyarrow as pa

# Pandas DataFrame stored as Parquet
@cache.asset(expire=3600)
def get_sales_data():
 return pd.DataFrame({'sales': [100, 200, 300], 'region': ['A', 'B', 'C']})

sales = get_sales_data() # Stored as .parquet file

# Polars DataFrame stored as Parquet 
@cache.asset(expire=3600)
def get_analytics_data():
 return pl.DataFrame({'metric': [1.1, 2.2, 3.3], 'date': ['2023-01', '2023-02', '2023-03']})

analytics = get_analytics_data()

# NumPy array stored as .npy file
@cache.asset(expire=3600)
def get_model_weights():
 return np.random.random((1000, 100))

weights = get_model_weights()

# PyArrow Table stored as Arrow IPC format
@cache.asset(expire=3600, format="arrow")
def get_arrow_table():
 return pa.table({'col1': [1, 2, 3], 'col2': ['a', 'b', 'c']})

table = get_arrow_table()

# JSON for dicts/lists (when JSON-serializable)
@cache.asset(expire=3600)
def get_config():
 return {"database": "postgres", "port": 5432, "features": ["auth", "cache"]}

config = get_config() # Stored as .json file

# Force specific format
@cache.asset(expire=3600, format="parquet")
def get_mixed_data():
 # Even if data could be JSON, force Parquet format
 return pl.DataFrame({'values': [1, 2, 3]})

Extensible Asset Handlers

Create custom handlers for new data formats:

class CSVHandler(AssetHandler):
 format_name = "csv"

 def can_handle(self, data):
 return hasattr(data, 'to_csv') # Pandas-like objects

 def save(self, data, path):
 data.to_csv(path, index=False)

 def load(self, path):
 import pandas as pd
 return pd.read_csv(path)

 def is_available(self):
 try:
 import pandas
 return True
 except ImportError:
 return False

# Register custom handler
cache.register_asset_handler(CSVHandler())

# Now DataFrames can be stored as CSV
@cache.asset(format="csv")
def get_report():
 return pd.DataFrame({'report': ['Q1', 'Q2'], 'revenue': [100000, 120000]})

Clearing Cache by Name or Function

# Clear all cache
cache.clear()

# Clear specific key
cache.clear("user:123")

# Clear specific function's cached results
cache.clear(expensive_function)

# Clear multiple items
cache.clear("key1", "key2", my_function)

# Clear with list
cache.clear(["session:abc", "temp:xyz"])

Storage Backends

# SQLite database for metadata and small values
cache = Cache(
 directory="/path/to/cache",
 dbname="my_cache.db",
 sqlite_timeout=60,
 min_file_size=32768 # Store values < 32KB in SQLite, larger as files
)

# Large files stored on disk with appropriate extensions
# Small values stored as BLOBs in SQLite database
# Automatic file/database decision based on size

Eviction Policies

# Least Recently Used (LRU)
cache = Cache(eviction_policy="least-recently-used", size_limit=1024*1024*1024) # 1GB

# Least Frequently Used (LFU) 
cache = Cache(eviction_policy="least-frequently-used")

# Least Recently Stored
cache = Cache(eviction_policy="least-recently-stored")

# No eviction
cache = Cache(eviction_policy="none")

Pickle Fallback

When native format handlers fail or aren’t available, automatic fallback to pickle:

# Custom object without specific handler
class CustomData:
 def __init__(self, value):
 self.value = value

@cache.asset(expire=3600)
def get_custom_object():
 return CustomData("important_data")

# Automatically falls back to pickle storage
obj = get_custom_object()

Advanced Configuration

cache = Cache(
 directory="/var/cache/myapp",
 dbname="cache.db",
 eviction_policy="least-recently-used",
 size_limit=2**30, # 1GB total cache size
 cull_limit=100, # Remove 100 items when culling
 min_file_size=2**15, # 32KB threshold for file vs database storage
 sqlite_cache_size=8192, # SQLite page cache
 sqlite_mmap_size=2**26, # 64MB memory mapping
 asset_handlers=[CustomHandler(), JSONHandler(), PickleHandler()]
)

Implementation code

This is a zero-dependency implementation for disk-based caching with native format support. It includes the Cache class, AssetHandler base class, and various asset handlers for JSON, Parquet, Arrow, Numpy, and Pickle.

import os
import io
import time
import errno
import shutil
import codecs
import pickle
import sqlite3
import hashlib
import tempfile
import threading
import contextlib
import typing as t
import pickletools
from pathlib import Path
from functools import partial, wraps

# Type variables for preserving function signatures
P = t.ParamSpec('P')
T = t.TypeVar('T')

EVICTION_POLICY = {
 "none": {"init": None, "get": None, "cull": None},
 "least-recently-stored": {
 "init": "CREATE INDEX IF NOT EXISTS Cache_store_time ON Cache (store_time)",
 "get": None,
 "cull": "SELECT {fields} FROM Cache ORDER BY store_time LIMIT ?",
 },
 "least-recently-used": {
 "init": "CREATE INDEX IF NOT EXISTS Cache_access_time ON Cache (access_time)",
 "get": "access_time = {now}",
 "cull": "SELECT {fields} FROM Cache ORDER BY access_time LIMIT ?",
 },
 "least-frequently-used": {
 "init": "CREATE INDEX IF NOT EXISTS Cache_access_count ON Cache (access_count)",
 "get": "access_count = access_count + 1",
 "cull": "SELECT {fields} FROM Cache ORDER BY access_count LIMIT ?",
 },
}
MODE_NONE = 0
MODE_RAW = 1
MODE_BINARY = 2
MODE_TEXT = 3
MODE_PICKLE = 4


class AssetHandler:
 """Base class for asset format handlers.

 Subclasses should override format_name, can_handle, save, and load methods.
 """

 format_name: str = "base"

 def can_handle(self, data: t.Any) -> bool:
 """Check if this handler can process the given data.

 Args:
 data: The data to check

 Returns:
 bool: True if this handler can process the data
 """
 raise NotImplementedError

 def save(self, data: t.Any, path: str) -> None:
 """Save data to file.

 Args:
 data: The data to save
 path: The file path to save to
 """
 raise NotImplementedError

 def load(self, path: str) -> t.Any:
 """Load data from file.

 Args:
 path: The file path to load from

 Returns:
 The loaded data
 """
 raise NotImplementedError

 def is_available(self) -> bool:
 """Check if required dependencies are available.

 Returns:
 bool: True if all dependencies are available
 """
 return True


class PickleHandler(AssetHandler):
 """Default pickle handler for any Python object."""

 format_name = "pickle"

 def can_handle(self, data: t.Any) -> bool:
 return True # Can handle anything

 def save(self, data: t.Any, path: str) -> None:
 with open(path, "wb") as f:
 pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)

 def load(self, path: str) -> t.Any:
 with open(path, "rb") as f:
 return pickle.load(f)


class JSONHandler(AssetHandler):
 """JSON handler for dicts and lists."""

 format_name = "json"

 def can_handle(self, data: t.Any) -> bool:
 if not isinstance(data, (dict, list)):
 return False
 try:
 import json

 json.dumps(data)
 return True
 except (ImportError, TypeError, ValueError):
 return False

 def save(self, data: t.Any, path: str) -> None:
 import json

 with open(path, "w", encoding="utf-8") as f:
 json.dump(data, f, default=str, sort_keys=False, separators=(",", ": "))

 def load(self, path: str) -> t.Any:
 import json

 with open(path, "r", encoding="utf-8") as f:
 return json.load(f)

 def is_available(self) -> bool:
 try:
 import json # noqa: F401

 return True
 except ImportError:
 return False


class PolarsHandler(AssetHandler):
 """Polars DataFrame handler."""

 format_name = "parquet"

 def can_handle(self, data: t.Any) -> bool:
 try:
 return (
 hasattr(data, "__class__")
 and data.__class__.__module__ == "polars.dataframe.frame"
 and data.__class__.__name__ == "DataFrame"
 )
 except Exception:
 return False

 def save(self, data: t.Any, path: str) -> None:
 data.write_parquet(path)

 def load(self, path: str) -> t.Any:
 import polars as pl

 return pl.read_parquet(path)

 def is_available(self) -> bool:
 try:
 import polars # noqa: F401

 return True
 except ImportError:
 return False


class PandasHandler(AssetHandler):
 """Pandas DataFrame handler."""

 format_name = "parquet"

 def can_handle(self, data: t.Any) -> bool:
 try:
 return (
 hasattr(data, "__class__")
 and data.__class__.__module__ == "pandas.core.frame"
 and data.__class__.__name__ == "DataFrame"
 )
 except Exception:
 return False

 def save(self, data: t.Any, path: str) -> None:
 data.to_parquet(path)

 def load(self, path: str) -> t.Any:
 import pandas as pd

 return pd.read_parquet(path)

 def is_available(self) -> bool:
 try:
 import pandas # noqa: F401

 return True
 except ImportError:
 return False


class NumpyHandler(AssetHandler):
 """NumPy array handler."""

 format_name = "numpy"

 def can_handle(self, data: t.Any) -> bool:
 try:
 return (
 hasattr(data, "__class__")
 and data.__class__.__module__ == "numpy"
 and data.__class__.__name__ == "ndarray"
 )
 except Exception:
 return False

 def save(self, data: t.Any, path: str) -> None:
 import numpy as np

 # Remove .npy extension if present since np.save adds it automatically
 if path.endswith('.npy'):
 path = path[:-4]
 np.save(path, data)

 def load(self, path: str) -> t.Any:
 import numpy as np

 # np.save automatically adds .npy, so we need to handle both cases
 if not path.endswith('.npy') and not os.path.exists(path):
 path = path + '.npy'
 return np.load(path)

 def is_available(self) -> bool:
 try:
 import numpy # noqa: F401

 return True
 except ImportError:
 return False


class ArrowHandler(AssetHandler):
 """PyArrow Table handler."""

 format_name = "arrow"

 def can_handle(self, data: t.Any) -> bool:
 try:
 # Check for PyArrow Table
 return (
 hasattr(data, "__class__")
 and data.__class__.__module__ == "pyarrow.lib"
 and data.__class__.__name__ == "Table"
 )
 except Exception:
 return False

 def save(self, data: t.Any, path: str) -> None:
 import pyarrow as pa

 # Save as Arrow IPC format (Arrow file)
 with pa.OSFile(path, "wb") as sink:
 with pa.ipc.new_file(sink, data.schema) as writer:
 writer.write_table(data)

 def load(self, path: str) -> t.Any:
 import pyarrow as pa

 # Read Arrow IPC format file
 with pa.memory_map(path) as source:
 return pa.ipc.open_file(source).read_all()

 def is_available(self) -> bool:
 try:
 import pyarrow # noqa: F401

 return True
 except ImportError:
 return False


def delete(path: str | Path) -> None:
 try:
 if os.path.isdir(path):
 shutil.rmtree(path, ignore_errors=True)
 else:
 if os.path.exists(path):
 os.remove(path)
 except Exception:
 pass


class _Constant(tuple):
 """Pretty display of immutable constant."""

 def __new__(cls, name):
 return tuple.__new__(cls, (name,))

 def __repr__(self):
 return "%s" % self[0]


ENOVAL = _Constant("ENOVAL")
UNKNOWN = _Constant("UNKNOWN")


def sqlite_execute_with_retry(conn: sqlite3.Connection, statement: str, parameters: t.Iterable = ()) -> sqlite3.Cursor:
 """Execute a SQL statement with retry on database lock.

 Re-try the statement if the error is "database is locked" for up to 60 seconds.

 Args:
 conn: SQLite connection.
 statement: SQL statement.
 parameters: SQL statement parameters.

 Returns:
 sqlite3.Cursor: a cursor object

 Raises:
 sqlite3.OperationalError: if the error is not "database is locked".
 TimeoutError: if the database is locked for more than 60 seconds.
 """
 start = time.time()
 while True:
 try:
 return conn.execute(statement, parameters) # type: ignore[no-untyped-call]
 except sqlite3.OperationalError as exc:
 if str(exc) != "database is locked":
 raise # re-raise the original exception
 diff = time.time() - start
 if diff > 60:
 raise TimeoutError("SQLite database is locked for more than 60 seconds") from None
 time.sleep(0.001)


def full_name(func):
 """Return full name of `func` by adding the module and function name."""
 return func.__module__ + "." + func.__qualname__


def args_to_key(base: tuple, args: tuple, kwargs: dict, typed: bool | None, ignore: tuple | None) -> tuple:
 """Create cache key out of function arguments.

 Args:
 base: base of key
 args: function arguments
 kwargs: function keyword arguments
 typed: include types in cache key
 ignore: positional or keyword args to ignore

 Returns:
 cache key tuple
 """
 ignore = ignore or ()
 args = tuple(arg for index, arg in enumerate(args) if index not in ignore)
 key = base + args + (None,)

 if kwargs:
 kwargs = {key: val for key, val in kwargs.items() if key not in ignore}
 sorted_items = sorted(kwargs.items())

 for item in sorted_items:
 key += item

 if typed:
 key += tuple(type(arg) for arg in args)

 if kwargs:
 key += tuple(type(value) for _, value in sorted_items)

 return key


class Cache:
 """Disk and file based caching.

 This class provides a cache that stores data on disk using SQLite as the backend.
 It supports various eviction policies, file storage, and transaction management.

 It can store data in different formats using asset handlers, allowing for efficient
 storage and retrieval of data assets such as JSON, Polars DataFrames, Pandas DataFrames,
 NumPy arrays, and Python objects using Pickle.

 Args:
 directory: Directory to store cache files.
 sqlite_timeout: SQLite connection timeout in seconds.
 dbname: SQLite database name.
 pickle_protocol: Pickle protocol version.
 eviction_policy: Eviction policy for the cache. Options are "none", "least-recently-stored",
 "least-recently-used", "least-frequently-used". Default is "least-recently-stored".
 cull_limit: Number of items to cull when the cache is full.
 size_limit: Maximum cache size in bytes.
 min_file_size: Minimum file size in bytes before storing as file. Otherwise store as blob in database.
 sqlite_mmap_size: SQLite mmap size in bytes.
 sqlite_cache_size: SQLite cache size in pages.
 asset_handlers: List of AssetHandler instances for handling different data formats.
 """

 def __init__(
 self,
 directory: str | None = None,
 sqlite_timeout: int = 60,
 dbname: str = "cache.db",
 pickle_protocol: int = pickle.HIGHEST_PROTOCOL,
 eviction_policy: t.Literal[
 "none",
 "least-recently-stored",
 "least-recently-used",
 "least-frequently-used",
 ] = "least-recently-stored",
 cull_limit: int = 10,
 size_limit: int = 2**30, # 1GB
 min_file_size: int = 2**15, # 32KB
 sqlite_mmap_size: int = 2**26, # 64MB
 sqlite_cache_size: int = 2**13, # 8,192 pages
 asset_handlers: t.List[AssetHandler] | None = None,
 ) -> None:
 if directory is None:
 directory = tempfile.mkdtemp(prefix="webcache-")
 directory = str(directory)
 directory = os.path.expanduser(directory)
 directory = os.path.expandvars(directory)
 self.directory = directory
 self.dbname = dbname
 self.sqlite_timeout = sqlite_timeout
 self._local = threading.local()
 self._txn_id = None
 self.pickle_protocol = pickle_protocol
 self.eviction_policy = eviction_policy
 self.cull_limit = cull_limit
 self.size_limit = size_limit
 self.min_file_size = min_file_size
 self.sqlite_mmap_size = sqlite_mmap_size
 self.sqlite_cache_size = sqlite_cache_size

 # Initialize asset handlers
 if asset_handlers is None:
 # Default handlers
 self.asset_handlers = [
 JSONHandler(),
 ArrowHandler(),
 PolarsHandler(),
 PandasHandler(),
 NumpyHandler(),
 PickleHandler(), # Pickle last as fallback
 ]
 else:
 self.asset_handlers = asset_handlers
 # Ensure PickleHandler is always available as fallback
 if not any(isinstance(h, PickleHandler) for h in self.asset_handlers):
 self.asset_handlers.append(PickleHandler())

 if not os.path.isdir(directory):
 try:
 os.makedirs(directory, 0o755)
 except OSError as error:
 if error.errno != errno.EEXIST:
 raise EnvironmentError(
 error.errno,
 'Cache directory "%s" does not exist and could not be created' % self.directory,
 ) from None

 con = self.connect()

 # Set sqlite wal journal mode, auto vacuum, and mmap size, synchronous, cache size
 con.execute("PRAGMA journal_mode = WAL")
 con.execute("PRAGMA auto_vacuum = FULL")
 con.execute("PRAGMA mmap_size = %d" % self.sqlite_mmap_size)
 con.execute("PRAGMA synchronous = NORMAL")
 con.execute("PRAGMA cache_size = %d" % self.sqlite_cache_size)

 ((self._page_size,),) = con.execute("PRAGMA page_size").fetchall()

 con.execute("CREATE TABLE IF NOT EXISTS Settings (key TEXT NOT NULL UNIQUE, value)")
 con.execute("INSERT OR REPLACE INTO Settings VALUES (?, ?)", ("size", 0))
 con.execute("""CREATE TABLE IF NOT EXISTS Cache (
 rowid INTEGER PRIMARY KEY,
 key BLOB,
 raw INTEGER,
 store_time REAL,
 expire_time REAL,
 access_time REAL,
 access_count INTEGER DEFAULT 0,
 size INTEGER DEFAULT 0,
 mode INTEGER DEFAULT 0,
 filename TEXT,
 value BLOB
 )""")
 con.execute("CREATE UNIQUE INDEX IF NOT EXISTS Cache_key_raw ON Cache(key, raw)")
 con.execute("CREATE INDEX IF NOT EXISTS Cache_expire_time ON Cache (expire_time) WHERE expire_time IS NOT NULL")

 # Use triggers to keep size metadata up to date
 con.execute("""CREATE TRIGGER IF NOT EXISTS Settings_size_insert
 AFTER INSERT ON Cache FOR EACH ROW BEGIN
 UPDATE Settings SET value = value + NEW.size
 WHERE key = "size"; END""")
 con.execute("""CREATE TRIGGER IF NOT EXISTS Settings_size_update
 AFTER UPDATE ON Cache FOR EACH ROW BEGIN
 UPDATE Settings
 SET value = value + NEW.size - OLD.size
 WHERE key = "size"; END""")
 con.execute("""CREATE TRIGGER IF NOT EXISTS Settings_size_delete
 AFTER DELETE ON Cache FOR EACH ROW BEGIN
 UPDATE Settings SET value = value - OLD.size
 WHERE key = "size"; END""")

 query = EVICTION_POLICY[self.eviction_policy]["init"]
 if query is not None:
 con.execute(query)

 def connect(self):
 local_pid = getattr(self._local, "pid", None)
 pid = os.getpid()

 if local_pid != pid:
 self.close()
 self._local.pid = pid

 con = getattr(self._local, "con", None)

 if con is None:
 con = self._local.con = sqlite3.connect(
 os.path.join(self.directory, self.dbname),
 timeout=self.sqlite_timeout,
 isolation_level=None,
 )

 return con

 def close(self):
 con: sqlite3.Connection = getattr(self._local, "con", None) # type: ignore[no-untyped-call]
 if con is None:
 return
 con.close()
 try:
 delattr(self._local, "con")
 except AttributeError:
 pass

 def _disk_remove(self, file_path):
 """Remove a file given by `file_path` with cross-thread and cross-process safety."""
 full_path = os.path.join(self.directory, file_path)
 full_dir, _ = os.path.split(full_path)
 with contextlib.suppress(OSError):
 os.remove(full_path)
 with contextlib.suppress(OSError):
 os.removedirs(full_dir)

 @contextlib.contextmanager
 def transact(self, retry=False, filename=None) -> t.Iterator[t.Tuple[sqlite3.Connection, t.Callable]]:
 """Transaction context manager locking the cache.

 Args:
 retry: whether to retry the transaction if it fails.
 filename: filename to remove if the transaction fails.

 Raises:
 CacheDatabaseTransactionTimeout: if the transaction times out.

 Example:
 Wrap a block of code in a transaction:

 >>> with cache.transact() as (con, _):
 ... con.execute("CREATE TABLE IF NOT EXISTS test (id INTEGER PRIMARY KEY, name TEXT)")
 ... con.execute("INSERT INTO test (name) VALUES (?)", ("Alice",))
 """
 con: sqlite3.Connection = self.connect()
 filenames = []
 _disk_remove = self._disk_remove
 tid = threading.get_ident()
 txn_id = self._txn_id

 if tid == txn_id:
 begin = False
 else:
 while True:
 try:
 con.execute("BEGIN IMMEDIATE")
 begin = True
 self._txn_id = tid
 break
 except sqlite3.OperationalError:
 if retry:
 continue
 if filename is not None:
 _disk_remove(filename)
 raise TimeoutError from None

 try:
 yield con, filenames.append
 except BaseException:
 if begin:
 assert self._txn_id == tid
 self._txn_id = None
 con.execute("ROLLBACK")
 raise
 else:
 if begin:
 assert self._txn_id == tid
 self._txn_id = None
 con.execute("COMMIT")
 for name in filenames:
 if name is not None:
 _disk_remove(name)

 def setting(self, key: str, value: t.Any = ENOVAL, update: bool = True) -> t.Any:
 """Get or set a setting in the cache."""
 con = self.connect()

 if value is ENOVAL:
 select = "SELECT value FROM Settings WHERE key = ?"
 ((value,),) = sqlite_execute_with_retry(con, select, (key,)).fetchall()
 return value

 if update:
 statement = "UPDATE Settings SET value = ? WHERE key = ?"
 sqlite_execute_with_retry(con, statement, (value, key))

 def volume(self) -> int:
 """Return estimated total size of cache on disk in bytes."""
 con = self.connect()
 ((page_count,),) = con.execute("PRAGMA page_count").fetchall()
 total_size = self._page_size * page_count + self.setting("size")
 return total_size

 def _disk_put(self, key):
 """Convert `key` to fields key and raw for Cache table.

 Args:
 key: key to be stored in cache.

 Returns:
 Tuple[sqlite3.Binary, bool]: a tuple of the key and a boolean indicating whether the key is a byte string.
 """
 # pylint: disable=unidiomatic-typecheck
 type_key = type(key)

 if type_key is bytes:
 return sqlite3.Binary(key), True
 elif (
 (type_key is str)
 or (type_key is int and -9223372036854775808 <= key <= 9223372036854775807)
 or (type_key is float)
 ):
 return key, True
 else:
 #
 data = pickle.dumps(key, protocol=self.pickle_protocol)
 result = pickletools.optimize(data)
 return sqlite3.Binary(result), False

 def touch(self, key: str, expire: float | None = None, retry: bool = False) -> bool:
 """Touch `key` in cache and update `expire` time."""
 now = time.time()
 db_key, raw = self._disk_put(key)
 expire_time = None if expire is None else now + expire

 with self.transact(retry) as (con, _):
 rows = con.execute(
 "SELECT rowid, expire_time FROM Cache WHERE key = ? AND raw = ?",
 (db_key, raw),
 ).fetchall()

 if rows:
 ((rowid, old_expire_time),) = rows

 if old_expire_time is None or old_expire_time > now:
 con.execute(
 "UPDATE Cache SET expire_time = ? WHERE rowid = ?",
 (expire_time, rowid),
 )
 return True

 return False

 def _disk_fetch(self, mode: int, filename: str, value: t.Any, read: bool):
 """Convert fields `mode`, `filename`, and `value` from Cache table to value.

 If mode is MODE_RAW, return value as bytes. If mode is MODE_BINARY and read is true, return value as file handle,
 otherwise return value as bytes. If mode is MODE_TEXT, return value as string. If mode is MODE_PICKLE, read value
 as pickle and return the result.

 Args:
 mode: mode of the value. Options are MODE_RAW, MODE_BINARY, MODE_TEXT, MODE_PICKLE.
 filename: filename of the value.
 value: value to be fetched.
 read: whether to read the value.

 Returns:
 Any: the fetched value as str, bytes, file handle, or any other type if mode is MODE_PICKLE.
 """
 if mode == MODE_RAW:
 return bytes(value) if type(value) is sqlite3.Binary else value
 elif mode == MODE_BINARY:
 if read:
 return open(os.path.join(self.directory, filename), "rb")
 else:
 with open(os.path.join(self.directory, filename), "rb") as reader:
 return reader.read()
 elif mode == MODE_TEXT:
 full_path = os.path.join(self.directory, filename)
 with open(full_path, "r", encoding="UTF-8") as reader:
 return reader.read()
 elif mode == MODE_PICKLE:
 if value is None:
 with open(os.path.join(self.directory, filename), "rb") as reader:
 return pickle.load(reader)
 else:
 return pickle.load(io.BytesIO(value))

 def get(
 self,
 key: str,
 default: t.Any = None,
 read: bool = False,
 return_expire_time: bool = False,
 retry: bool = False,
 ):
 """Get `key` from cache.

 Args:
 key: key to be retrieved from cache.
 default: default value to return if key is not found.
 read: return file handle to value.
 expire_time: return expire time.
 retry: whether to retry on database lock.
 """
 db_key, raw = self._disk_put(key)
 update_column: str = EVICTION_POLICY[self.eviction_policy]["get"]
 select = """SELECT rowid, expire_time, mode, filename, value
 FROM Cache WHERE key = ? AND raw = ?
 AND (expire_time IS NULL OR expire_time > ?)
 """

 if return_expire_time:
 default = (default, None)

 with self.transact(retry) as (con, _):
 rows = con.execute(select, (db_key, raw, time.time())).fetchall()
 if not rows:
 return default
 ((rowid, db_expire_time, mode, filename, db_value),) = rows
 try:
 value = self._disk_fetch(mode, filename, db_value, read)
 except IOError:
 # Key was deleted before we could retrieve result.
 return default

 if update_column is not None:
 now = time.time()
 update = "UPDATE Cache SET %s WHERE rowid = ?"
 con.execute(update % update_column.format(now=now), (rowid,))

 if return_expire_time:
 return value, db_expire_time
 return value

 def __getitem__(self, key):
 """Return corresponding value for `key` from cache."""
 value = self.get(key, default=ENOVAL, retry=True)
 if value is ENOVAL:
 raise KeyError(key)
 return value

 def read(self, key, retry=False):
 """Return file handle value corresponding to `key` from cache."""
 handle = self.get(key, default=ENOVAL, read=True, retry=retry)
 if handle is ENOVAL:
 raise KeyError(key)
 return handle

 def __contains__(self, key: str) -> bool:
 """Return `True` if `key` matching item is found in cache."""
 con = self.connect()
 db_key, raw = self._disk_put(key)
 select = "SELECT rowid FROM Cache WHERE key = ? AND raw = ? AND (expire_time IS NULL OR expire_time > ?)"
 rows = con.execute(select, (db_key, raw, time.time())).fetchall()
 return bool(rows)

 def exists(self, key: str) -> bool:
 """Return `True` if `key` matching item is found in cache."""
 return key in self

 def _disk_filename(self, key: t.Any = UNKNOWN, value: t.Any = UNKNOWN):
 """Return filename and full-path tuple for file storage."""
 hex_name = codecs.encode(os.urandom(16), "hex").decode("utf-8")
 sub_dir = os.path.join(hex_name[:2], hex_name[2:4])
 name = hex_name[4:] + ".val"
 filename = os.path.join(sub_dir, name)
 full_path = os.path.join(self.directory, filename)
 return filename, full_path

 def _safe_filename(self, base_filename: str, extension: str, max_length: int = 255) -> str:
 """Create a safe filename that doesn't exceed filesystem limits.

 Args:
 base_filename: Base filename without extension
 extension: File extension (without dot)
 max_length: Maximum filename length (default 255 for most filesystems)

 Returns:
 Safe filename that fits within length limits
 """
 full_name = f"{base_filename}.{extension}"

 # If filename is already short enough, return as-is
 if len(full_name) <= max_length:
 return full_name

 # Calculate how much we need to truncate
 # Reserve space for extension and a hash separator
 available_length = max_length - len(extension) - 1 - 8 # 8 chars for hash

 if available_length <= 0:
 # Extension is too long, use only hash

 hash_str = hashlib.md5(base_filename.encode()).hexdigest()[:8]
 return f"{hash_str}.{extension}"

 # Truncate base and add hash to maintain uniqueness

 truncated_base = base_filename[:available_length]
 hash_str = hashlib.md5(base_filename.encode()).hexdigest()[:8]
 return f"{truncated_base}_{hash_str}.{extension}"

 def _disk_store(self, value, read, key=UNKNOWN):
 """Convert `value` to fields size, mode, filename, and value for Cache table.

 Args:
 value: value to convert
 read: True when value is file-like object
 key: key for item (default UNKNOWN)

 Returns:
 (size, mode, filename, value) tuple for Cache table
 """
 # pylint: disable=unidiomatic-typecheck
 type_value = type(value)
 min_file_size = self.min_file_size

 if (
 (type_value is str and len(value) < min_file_size)
 or (type_value is int and -9223372036854775808 <= value <= 9223372036854775807)
 or (type_value is float)
 ):
 return 0, MODE_RAW, None, value
 elif type_value is bytes:
 if len(value) < min_file_size:
 return 0, MODE_RAW, None, sqlite3.Binary(value)
 else:
 filename, full_path = self._disk_filename(key, value)
 self._disk_write(full_path, io.BytesIO(value), "xb")
 return len(value), MODE_BINARY, filename, None
 elif type_value is str:
 filename, full_path = self._disk_filename(key, value)
 self._disk_write(full_path, io.StringIO(value), "x", "UTF-8")
 size = os.path.getsize(full_path)
 return size, MODE_TEXT, filename, None
 elif read:
 reader = partial(value.read, 2**22)
 filename, full_path = self._disk_filename(key, value)
 iterator = iter(reader, b"")
 size = self._disk_write(full_path, iterator, "xb")
 return size, MODE_BINARY, filename, None
 else:
 result = pickle.dumps(value, protocol=self.pickle_protocol)

 if len(result) < min_file_size:
 return 0, MODE_PICKLE, None, sqlite3.Binary(result)
 else:
 filename, full_path = self._disk_filename(key, value)
 self._disk_write(full_path, io.BytesIO(result), "xb")
 return len(result), MODE_PICKLE, filename, None

 def _disk_write(self, full_path, iterator, mode, encoding=None):
 full_dir, _ = os.path.split(full_path)
 for count in range(1, 11):
 try:
 # Ensure directory exists - use exist_ok to handle race conditions
 os.makedirs(full_dir, exist_ok=True)

 # Try to open the file - if directory was deleted, this will fail
 with open(full_path, mode, encoding=encoding) as writer:
 size = 0
 for chunk in iterator:
 size += len(chunk)
 writer.write(chunk)
 return size

 except (OSError, IOError) as e:
 # Handle various filesystem errors including permission issues,
 # directory deletion, disk full, etc.
 if count == 10:
 # Give up after 10 tries
 raise OSError(f"Failed to write file after {count} attempts: {e}") from e

 # Clean up partial file if it exists
 with contextlib.suppress(OSError):
 if os.path.exists(full_path):
 os.remove(full_path)

 # Brief delay before retry to allow for transient conditions to resolve
 time.sleep(0.001 * count) # Exponential backoff
 continue

 def set(self, key, value, expire=None, read=False, retry=False):
 """Set corresponding `value` for `key` in cache

 Args:
 key: key name
 value: value to store
 expire: expire time in seconds.
 read: whether to read the value.
 retry: whether to retry on database lock.
 """
 now = time.time()
 db_key, raw = self._disk_put(key)
 expire_time = None if expire is None else now + expire
 size, mode, filename, db_value = self._disk_store(value, read, key=key)
 columns = (expire_time, size, mode, filename, db_value)

 with self.transact(retry, filename) as (con, cleanup):
 rows = con.execute(
 "SELECT rowid, filename FROM Cache WHERE key = ? AND raw = ?",
 (db_key, raw),
 ).fetchall()
 if rows:
 ((rowid, old_filename),) = rows
 cleanup(old_filename)
 self._row_update(rowid, now, columns)
 else:
 self._row_insert(db_key, raw, now, columns)
 self._cull(now, con, cleanup)
 return True

 def __setitem__(self, key, value):
 """Set corresponding `value` for `key` in cache."""
 self.set(key, value, retry=True)

 def _row_update(self, rowid, now, columns):
 con = self.connect()
 expire_time, size, mode, filename, value = columns
 con.execute(
 """UPDATE Cache SET
 store_time = ?,
 expire_time = ?,
 access_time = ?,
 access_count = ?,
 size = ?,
 mode = ?,
 filename = ?,
 value = ?
 WHERE rowid = ?
 """,
 (
 now, # store_time
 expire_time,
 now, # access_time
 0, # access_count
 size,
 mode,
 filename,
 value,
 rowid,
 ),
 )

 def _row_insert(self, key, raw, now, columns):
 con = self.connect()
 expire_time, size, mode, filename, value = columns
 con.execute(
 """INSERT INTO Cache(
 key, raw, store_time, expire_time, access_time,
 access_count, size, mode, filename, value
 ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
 """,
 (
 key,
 raw,
 now, # store_time
 expire_time,
 now, # access_time
 0, # access_count
 size,
 mode,
 filename,
 value,
 ),
 )

 def _cull(self, now: int | float, con: sqlite3.Connection, cleanup: t.Callable, limit: int | None = None):
 cull_limit = self.cull_limit if limit is None else limit
 if cull_limit == 0:
 return

 # Evict expired keys
 select_expired_template = (
 "SELECT %s FROM Cache WHERE expire_time IS NOT NULL AND expire_time < ? ORDER BY expire_time LIMIT ?"
 )
 select_expired = select_expired_template % "filename"
 rows = con.execute(select_expired, (now, cull_limit)).fetchall()
 if rows:
 delete_expired = "DELETE FROM Cache WHERE rowid IN (%s)" % (select_expired_template % "rowid")
 con.execute(delete_expired, (now, cull_limit))
 for (filename,) in rows:
 cleanup(filename)
 cull_limit -= len(rows)
 if cull_limit == 0:
 return

 # Evict keys by policy
 select_policy = EVICTION_POLICY[self.eviction_policy]["cull"]
 if select_policy is None or self.volume() < self.size_limit:
 return
 select_filename = select_policy.format(fields="filename", now=now)
 rows = con.execute(select_filename, (cull_limit,)).fetchall()
 if rows:
 delete = "DELETE FROM Cache WHERE rowid IN (%s)" % (select_policy.format(fields="rowid", now=now))
 con.execute(delete, (cull_limit,))
 for (filename,) in rows:
 cleanup(filename)

 def clear(self, *args, retry=False):
 """Remove items from cache.

 Args:
 *args: Optional arguments to specify what to clear:
 - No args: Clear all items (default behavior)
 - str: Clear items matching the key name
 - callable: Clear items for memoized/asset decorated functions
 - list/tuple: Clear items matching multiple keys/functions
 retry: Whether to retry on database lock

 Returns:
 int: Number of items cleared

 Examples:
 >>> cache = Cache()
 >>> cache.clear() # Clear all items
 >>> cache.clear("my_key") # Clear specific key
 >>> cache.clear(my_func) # Clear memoized function cache
 >>> cache.clear("key1", "key2", my_func) # Clear multiple items
 """
 if not args:
 # Clear all items (original behavior)
 select = "SELECT rowid, filename FROM Cache WHERE rowid > ? ORDER BY rowid LIMIT ?"
 select_args = [0, 100]
 return self._select_delete(select, select_args, retry=retry)

 # Build list of keys to clear
 keys_to_clear = []

 for arg in args:
 if isinstance(arg, str):
 # Direct key name
 keys_to_clear.append(arg)
 elif callable(arg):
 # Memoized or asset decorated function
 if hasattr(arg, "__cache_key__"):
 # This is a decorated function - we need to clear all its cached results
 # We'll use a pattern match on the function's full name
 func_name = full_name(arg)
 keys_to_clear.append(func_name)
 else:
 # Regular function, use its full name
 keys_to_clear.append(full_name(arg))
 elif isinstance(arg, (list, tuple)):
 # Recursively process lists/tuples
 for item in arg:
 if isinstance(item, str):
 keys_to_clear.append(item)
 elif callable(item):
 if hasattr(item, "__cache_key__"):
 keys_to_clear.append(full_name(item))
 else:
 keys_to_clear.append(full_name(item))
 else:
 # Convert other types to string keys
 keys_to_clear.append(str(arg))

 if not keys_to_clear:
 return 0

 return self._clear_specific_keys(keys_to_clear, retry=retry)

 def _clear_specific_keys(self, keys: t.List[str], retry: bool = False) -> int:
 """Clear cache entries for specific keys or key patterns.

 Args:
 keys: List of keys or function names to clear
 retry: Whether to retry on database lock

 Returns:
 Number of items cleared
 """
 total_cleared = 0

 for key_pattern in keys:
 # Handle exact key matches
 db_key, raw = self._disk_put(key_pattern)

 # Clear exact matches
 select = "SELECT rowid, filename FROM Cache WHERE key = ? AND raw = ?"
 try:
 with self.transact(retry) as (con, cleanup):
 rows = con.execute(select, (db_key, raw)).fetchall()
 if rows:
 rowids = [str(row[0]) for row in rows]
 delete = f"DELETE FROM Cache WHERE rowid IN ({','.join(rowids)})"
 con.execute(delete)

 for row in rows:
 cleanup(row[1]) # Clean up file

 total_cleared += len(rows)
 except TimeoutError:
 pass

 # Also clear function-based cache entries (for memoized/asset functions)
 # These have keys that start with the function name tuple
 try:
 with self.transact(retry) as (con, cleanup):
 # Get all non-raw cache entries and check them
 select = "SELECT rowid, filename, key FROM Cache WHERE raw = 0"
 rows = con.execute(select).fetchall()

 matching_rows = []
 for rowid, filename, key_blob in rows:
 try:
 # Deserialize the key to check if it matches our function
 key_tuple = pickle.loads(bytes(key_blob))
 if isinstance(key_tuple, tuple) and len(key_tuple) > 0:
 if key_tuple[0] == key_pattern:
 matching_rows.append((rowid, filename))
 except (pickle.PickleError, TypeError, IndexError):
 continue

 if matching_rows:
 rowids = [str(row[0]) for row in matching_rows]
 delete = f"DELETE FROM Cache WHERE rowid IN ({','.join(rowids)})"
 con.execute(delete)

 for _, filename in matching_rows:
 cleanup(filename)

 total_cleared += len(matching_rows)

 except (TimeoutError, sqlite3.Error):
 pass

 return total_cleared

 def _select_delete(self, select, args, row_index=0, arg_index=0, retry=False):
 count = 0
 delete = "DELETE FROM Cache WHERE rowid IN (%s)"
 try:
 while True:
 with self.transact(retry) as (con, cleanup):
 rows = con.execute(select, args).fetchall()

 if not rows:
 break

 count += len(rows)
 con.execute(delete % ",".join(str(row[0]) for row in rows))

 for row in rows:
 args[arg_index] = row[row_index]
 cleanup(row[-1])
 except TimeoutError:
 raise TimeoutError(count) from None
 return count

 def cleanup(self, force: bool = False, cache_timeout: int = 3600):
 """Remove all cache files that are older than `cache_timeout` seconds.

 Args:
 force (bool): If True, delete all cache files regardless of age.
 cache_timeout (int): Cache timeout in seconds.
 """
 for path in Path(self.directory).iterdir():
 if force:
 delete(Path(path))
 continue
 age = time.time() - os.stat(path).st_mtime
 if age > cache_timeout:
 delete(Path(path))

 def memoize(self, name=None, typed=False, expire=None, ignore=()) -> t.Callable[[t.Callable[P, T]], t.Callable[P, T]]:
 """Memoizing cache decorator.

 Decorator to wrap callable with memoizing function using cache.
 Repeated calls with the same arguments will lookup result in cache and
 avoid function evaluation.

 If name is set to None (default), the callable name will be determined
 automatically.

 When expire is set to zero, function results will not be set in the
 cache. Cache lookups still occur, however. Read
 :doc:`case-study-landing-page-caching` for example usage.

 If typed is set to True, function arguments of different types will be
 cached separately. For example, f(3) and f(3.0) will be treated as
 distinct calls with distinct results.

 The original underlying function is accessible through the __wrapped__
 attribute. This is useful for introspection, for bypassing the cache,
 or for rewrapping the function with a different cache.

 >>> cache = Cache()
 >>> @cache.memoize(expire=1, tag='fib')
 ... def fibonacci(number):
 ... if number == 0:
 ... return 0
 ... elif number == 1:
 ... return 1
 ... else:
 ... return fibonacci(number - 1) + fibonacci(number - 2)
 >>> print(fibonacci(100))
 354224848179261915075

 An additional `__cache_key__` attribute can be used to generate the
 cache key used for the given arguments.

 >>> key = fibonacci.__cache_key__(100)
 >>> print(cache[key])
 354224848179261915075

 Remember to call memoize when decorating a callable. If you forget,
 then a TypeError will occur. Note the lack of parenthenses after
 memoize below:

 >>> @cache.memoize
 ... def test():
 ... pass
 Traceback (most recent call last):
 ...
 TypeError: name cannot be callable

 Args:
 cache: cache to store callable arguments and return values
 str name: name given for callable (default None, automatic)
 bool typed: cache different types separately (default False)
 float expire: seconds until arguments expire
 (default None, no expiry)
 str tag: text to associate with arguments (default None)
 set ignore: positional or keyword args to ignore (default ())

 Returns:
 callable decorator

 """
 if callable(name):
 raise TypeError("name cannot be callable")

 def decorator(func: t.Callable[P, T]) -> t.Callable[P, T]:
 """Decorator created by memoize() for callable `func`."""
 base = (full_name(func),) if name is None else (name,)

 @wraps(func)
 def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
 """Wrapper for callable to cache arguments and return values."""
 key = wrapper.__cache_key__(*args, **kwargs) # type: ignore[no-untyped-call]
 result = self.get(key, default=ENOVAL, retry=True)

 if result is ENOVAL:
 result = func(*args, **kwargs)
 if expire is None or expire > 0:
 self.set(key, result, expire, retry=True)
 return result

 return t.cast(T, result)

 def __cache_key__(*args, **kwargs):
 """Make key for cache given function arguments."""
 return args_to_key(base, args, kwargs, typed, ignore)

 wrapper.__cache_key__ = __cache_key__ # type: ignore[no-untyped-call]
 return t.cast(t.Callable[P, T], wrapper)

 return decorator

 def register_asset_handler(self, handler: AssetHandler, prepend: bool = False) -> None:
 """Register a new asset handler.

 Args:
 handler: The AssetHandler instance to register
 prepend: If True, add to beginning of list (higher priority)
 """
 if prepend:
 self.asset_handlers.insert(0, handler)
 else:
 # Insert before PickleHandler if it exists
 pickle_idx = None
 for i, h in enumerate(self.asset_handlers):
 if isinstance(h, PickleHandler):
 pickle_idx = i
 break

 if pickle_idx is not None:
 self.asset_handlers.insert(pickle_idx, handler)
 else:
 self.asset_handlers.append(handler)

 def _get_asset_handler(self, data: t.Any, format: str | None = None) -> AssetHandler:
 """Get the appropriate asset handler for the data.

 Args:
 data: The data to find a handler for
 format: Optional format hint

 Returns:
 AssetHandler: The handler that can process this data
 """
 # If format is specified, try to find handler with that format name
 if format:
 for handler in self.asset_handlers:
 if handler.format_name == format and handler.is_available():
 return handler

 # Otherwise, find first handler that can handle the data
 for handler in self.asset_handlers:
 if handler.is_available() and handler.can_handle(data):
 return handler

 # This should never happen if PickleHandler is in the list
 raise ValueError("No handler found for data type")

 def _store_asset(self, key: str, data: t.Any, expire: float | None, format: str | None = None) -> None:
 """Store an asset using the appropriate handler.

 Args:
 key: Cache key
 data: Data to store
 expire: Expiration time in seconds
 format: Optional format hint
 """
 handler = self._get_asset_handler(data, format)

 # Generate filename with format extension and data type info to avoid collisions
 # Include data type information in the filename to prevent collisions
 data_type_info = f"{type(data).__module__}.{type(data).__name__}"
 # Create a hash of the data type to keep filename manageable
 type_hash = hashlib.md5(data_type_info.encode()).hexdigest()[:8]
 filename, full_path = self._disk_filename(key, str(data)[:100]) # Use truncated str for filename

 # Use safe filename to ensure it doesn't exceed filesystem limits
 # Only apply to the actual filename part, not the directory path
 dir_part, name_part = os.path.split(filename)
 name_base, _ = name_part.rsplit(".", 1)
 safe_name = self._safe_filename(f"{name_base}_{type_hash}", handler.format_name)
 filename = os.path.join(dir_part, safe_name)

 dir_part, name_part = os.path.split(full_path)
 full_path = os.path.join(dir_part, safe_name)

 # Save using handler
 full_dir, _ = os.path.split(full_path)
 os.makedirs(full_dir, exist_ok=True)

 try:
 handler.save(data, full_path)
 # For numpy handler, the actual file created might have .npy extension
 actual_path = full_path
 if isinstance(handler, NumpyHandler) and not os.path.exists(full_path):
 actual_path = full_path + '.npy'
 size = os.path.getsize(actual_path)
 except Exception as e:
 # Clean up on failure - check both possible paths
 for path_to_clean in [full_path, full_path + '.npy']:
 if os.path.exists(path_to_clean):
 os.remove(path_to_clean)
 raise e

 # Store metadata in database
 now = time.time()
 db_key, raw = self._disk_put(key)
 expire_time = None if expire is None else now + expire

 # Update filename to reflect the actual file created (for numpy handler)
 if isinstance(handler, NumpyHandler) and not os.path.exists(os.path.join(self.directory, filename)):
 filename = filename + '.npy'

 with self.transact(retry=True, filename=filename) as (con, cleanup):
 # Check if key already exists
 rows = con.execute(
 "SELECT rowid, filename FROM Cache WHERE key = ? AND raw = ?",
 (db_key, raw),
 ).fetchall()

 if rows:
 ((rowid, old_filename),) = rows
 cleanup(old_filename)
 # Update existing entry
 con.execute(
 """UPDATE Cache SET
 store_time = ?, expire_time = ?, access_time = ?, access_count = ?,
 size = ?, mode = ?, filename = ?, value = ?
 WHERE rowid = ?""",
 (now, expire_time, now, 0, size, MODE_BINARY, filename, None, rowid),
 )
 else:
 # Insert new entry
 con.execute(
 """INSERT INTO Cache(
 key, raw, store_time, expire_time, access_time,
 access_count, size, mode, filename, value
 ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
 (db_key, raw, now, expire_time, now, 0, size, MODE_BINARY, filename, None),
 )

 self._cull(now, con, cleanup)

 def _load_asset(self, path: str, format: str | None = None) -> t.Any:
 """Load an asset from file.

 Args:
 path: File path
 format: Optional format hint

 Returns:
 The loaded data
 """
 # Check if file exists
 if not os.path.exists(path):
 raise FileNotFoundError(f"Asset file not found: {path}")

 # Determine format from file extension if not provided
 if format is None:
 _, ext = os.path.splitext(path)
 format = ext.lstrip(".")

 # Try to load with the specified/detected format handler
 format_handler = None
 for handler in self.asset_handlers:
 if handler.format_name == format and handler.is_available():
 format_handler = handler
 break

 if format_handler:
 try:
 return format_handler.load(path)
 except Exception as e:
 # If specific format handler fails, try fallback strategies
 if not isinstance(format_handler, PickleHandler):
 # Try pickle fallback if the original format wasn't pickle
 pickle_handler = None
 for handler in self.asset_handlers:
 if isinstance(handler, PickleHandler):
 pickle_handler = handler
 break

 if pickle_handler:
 try:
 return pickle_handler.load(path)
 except Exception:
 # If pickle also fails, raise the original error
 raise ValueError(f"Failed to load asset with {format} format: {e}") from e

 # If no fallback or fallback failed, raise original error
 raise ValueError(f"Failed to load asset with {format} format: {e}") from e

 # No handler found for the format - try pickle as last resort
 pickle_handler = None
 for handler in self.asset_handlers:
 if isinstance(handler, PickleHandler):
 pickle_handler = handler
 break

 if pickle_handler:
 try:
 return pickle_handler.load(path)
 except Exception as e:
 raise ValueError(f"No handler found for format '{format}' and pickle fallback failed: {e}") from e

 raise ValueError(f"No handler found for format: {format}")

 def asset(
 self,
 name: str | None = None,
 typed: bool | None = False,
 expire: int | float | None = None,
 ignore: tuple = (),
 format=None,
 ) -> t.Callable[[t.Callable[P, T]], t.Callable[P, T]]:
 """Asset caching decorator.

 Decorator to wrap callable with asset caching function.
 Similar to memoize but stores data in native formats instead of pickle.

 Supports various data types through duck typing:
 - PyArrow Tables: stored as Arrow IPC files
 - Polars DataFrames: stored as Parquet files
 - Pandas DataFrames: stored as Parquet files
 - NumPy arrays: stored as .npy files
 - Dicts/Lists: stored as JSON or pickle based on content
 - Other types: fallback to pickle

 Args:
 name: name given for callable (default None, automatic)
 typed: cache different types separately (default False)
 expire: seconds until arguments expire (default None, no expiry)
 ignore: positional or keyword args to ignore (default ())
 format: force specific format (arrow, parquet, numpy, json, pickle)

 Returns:
 callable decorator
 """
 if callable(name):
 raise TypeError("name cannot be callable")

 def decorator(func: t.Callable[P, T]) -> t.Callable[P, T]:
 """Decorator created by asset() for callable `func`."""
 base = (full_name(func),) if name is None else (name,)

 @wraps(func)
 def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
 """Wrapper for callable to cache asset arguments and return values."""
 key = wrapper.__cache_key__(*args, **kwargs) # type: ignore[attr-defined]

 # Check if asset exists in cache
 con = self.connect()
 db_key, raw = self._disk_put(key)
 select = """SELECT filename, mode FROM Cache WHERE key = ? AND raw = ?
 AND (expire_time IS NULL OR expire_time > ?)"""

 rows = con.execute(select, (db_key, raw, time.time())).fetchall()
 if rows:
 ((filename, _),) = rows # mode not needed here, renamed to _
 if filename:
 # Load from file using appropriate loader
 full_path = os.path.join(self.directory, filename)
 return t.cast(T, self._load_asset(full_path, format))

 # Asset not in cache, compute it
 result = func(*args, **kwargs)

 if expire is None or expire > 0:
 # Store asset in appropriate format
 self._store_asset(key, result, expire, format)

 return result

 def __cache_key__(*args, **kwargs):
 """Make key for cache given function arguments."""
 return args_to_key(base, args, kwargs, typed, ignore)

 wrapper.__cache_key__ = __cache_key__ # type: ignore[attr-defined]
 return t.cast(t.Callable[P, T], wrapper)

 return decorator

Document Chunking for LLMs

Tue, 17 Jun 2025 00:00:00 +0000

I’ve found that splitting large documents into smaller chunks requires some trial and error to find the right strategy.

There are basically three competing factors:

Semantic coherence within chunks - the chunk should contain related information. A small chunk size loses context but is more coherent while a large chunk size is less coherent but contains more context.
Semantic separation between chunks - the chunks should be distinct enough to avoid overlap but not so distinct that they lose context. This means we want to avoid splitting chunks at arbitrary points like sentences or paragraphs and instead split them at semantic boundaries of sections and topics. Even having some overlap between chunks can be helpful to keep context.
Information preservation - chunks should be self-contained enough to be useful on their own. This means that we need to split text at context boundaries and ensure that the similar concepts are grouped together in the same chunk and not split across chunks.

Ultimately, the appropriate startegy depends on the type of document and the intended use case. For example, a user guide is organized into chapters and sections. The beginning of a chapter introduces the topic, followed by sections that provide detailed workflows and instructions. In this case, it makes sense to keep the workflows and instructions together in the same chunk. We want to avoid splitting a step in a workflow into two independent chunks that loses it’s place in the workflow and more importantly, the context of the step. If the document is more structured, like an invoice or technical specifications with tables, than a more sophisticated chunking strategy may be needed to preserve the relationships between the data points. This article focuses on chunking textual documents, such as reports and manuals, where the text is more free-form and less structured.

Table of contents:

Algorithm Overview
Example: Berkshire Hathaway Annual Report
Implementation Code

Algorithm Overview

Lets build a chunking algorithm based on these principles from the ground up. Conceptually, we want to start with sentences since they represent the basic unit of a coherent meaning. And from there we need a strategy for iteratively merging sentences that achieves a balance between the three stated objectives: semantic coherence within chunks, semantic separation between chunks, and information preservation.

So the algorithm consists of two main stages:

Semantic chunking breaks the document into semantically coherent chunks by:
- Splitting into sentences
- Computing embeddings for each sentence
- Using cosine distance between consecutive sentence embeddings to detect topic shifts
- Enforcing size constraints while adding overlap
Section grouping organizes chunks into a hierarchical structure by:
- Generating descriptive metadata (title and category) for each chunk using an LLM
- Computing weighted similarity between chunk metadata
- Grouping similar chunks into sections based on similarity thresholds
- Creating subsections within sections for fine-grained organization

Example: Berkshire Hathaway Annual Report

As an example, I’ve taken the 2024 Berkshire Hathaway Annual Report and passed it through the chunking process. First I generates sentences from the document using regex and then computes embeddings and uses cosine similarity to merge sentences into partially overlapping chunks. Next I take each chunk and use an LLM to generate metadata including a title and category for each chunk. Notice how there is consistency in the titles and categories across chunks. This is because the LLM is prompted with the trailing 5 metadata entries to provide context for the current chunk. Finally, I use the embedding approach to merge chunks into sections based on the semantic similarity of the title and category of each chunk. A tunable similarity threshold of 0.75 is used to combine chunks into sections. The resulting sections are more coherent and self-contained, making them easier to use in a RAG system.

The table below shows the title and category generated by an LLM for each chunk. Splits happen when the chunk similarity between the current chunk and the previous chunk is below a threshold of 0.75 or when the chunk size exceeds a maximum character limit. When there is a split, the table below reports the similarity of the current chunk to the previous chunk.

In the beginning of the letter to the shareholders, the chunks are smaller and more granular because Buffet is introducing several concepts and providing an overview of the risks, management philosophy, and strategic focus. Chunks are smaller, spanning a few sentences and the similarity between chunks is lower. As the letter progresses, the chunks become larger and more coherent, with higher similarity scores as Buffet discusses specific topics in more detail. The final sections of the letter are larger and more comprehensive, covering specific companies, investments, and their performance in depth. Notice that even though the similarity between chunks is high and the title and category are the same, the chunk size grows larger and a new split is created.

Chunk	Title	Category	Previous Chunk Similarity
1	Berkshire Hathaway Annual Report	Financial Report
2	Berkshire Hathaway Annual Report	Financial Report
3	Company Transparency & Reporting	Corporate Communication & Reporting	0.541
4	Responsibility & Communication	Corporate Communication & Reporting	0.74
5	Strategic Review & Ownership Dialogue	Corporate Communication & Reporting	0.696
6	Communication Strategy	Business & Investor Relations	0.576
7	Berkshire Hathaway’s Risk Management Approach	Financial Risk Management	0.479
8	Mistakes and Strategic Assessment	Business Strategy	0.528
9	Mistakes in Berkshire Acquisitions	Business Risk & Management	0.641
10	Mistakes in Hiring, Impacting Berkshire	Corporate Management & Risk
11	Mistakes in Hiring Assessment	Personnel Management & Decision Making	0.668
12	Painful Mistakes, Diminishing Returns	Financial & Strategic	0.506
13	Mistakes and Delay	Strategic & Operational	0.679
14	Mistakes and Analysis	Business Strategy	0.643
15	Mistakes and Their Impact	Business Strategy
16	Word Frequency Analysis	Text Analysis	0.37
17	Company Observations	Business & Communication	0.405
18	Behavioral Observations	Business & Risk	0.652
19	CEO Succession & Risk	Corporate Management & Risk	0.568
20	CEO Transition & Risk	Corporate Management & Risk
21	CEO Succession & Berkshire’s Risk	Corporate Strategy & Risk
22	Berkshire CEO Philosophy	Corporate Strategy & Leadership
23	Pete Liegl’s Legacy	Business & Leadership	0.53
24	Pete Liegl - A Wealthing Story	Business & Financial
25	Pete - Forest River Founder	Business & Founding	0.572
26	Forest River Acquisition	Business & Financial	0.638
27	RV Deal - Initial Communication	Business & Communication & Financial	0.648
28	Berkshire Acquisition Deal	Business & Finance	0.646
29	Meeting Details & Price Discussion	Business & Communication & Strategy	0.513
30	Meeting & Deal	Business & Communication
31	Business Meeting & Financial Planning	Business & Strategy	0.664
32	Berkshire Hathaway Business Deal	Business & Financial	0.576
33	Real Estate Deal	Business & Financial
34	Real Estate Lease Dispute	Business & Financial
35	Meeting Dynamics	Business & Communication & Strategy	0.514
36	Compensation Structure	Financial & Human Resources	0.429
37	Compensation Structure	Financial & Strategy
38	Berkshire’s Compensation Offer	Financial & Business	0.669
39	Berkshire’s Financial Strategy	Financial Strategy & Risk	0.677
40	Berkshire’s Early Success	Business & Financial	0.741
41	Simple Success	Business & Financial	0.737
42	Berkshire’s Strategic Mistakes	Business & Risk & Strategy	0.492
43	Berkshire’s Performance	Financial & Strategic	0.736
44	Mistakes in Berkshire Acquisitions	Business Risk & Management	0.645
45	Berkshire’s Strategic Imperfections	Business Strategy & Risk
46	Strategic Focus & Partnership	Business & Strategy	0.624
47	Strategic Imperative	Business & Strategy
48	CEO Mistakes & Analysis	Business Strategy & Risk	0.614
49	Mistakes in Berkshire Acquisitions	Business Risk & Management	0.745
50	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
51	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
52	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
53	Berkshire’s Mistakes & Strategic Challenges	Business & Risk & Strategy
54	Mistakes in Berkshire Acquisitions	Business Risk & Management
55	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
56	Berkshire’s Strategic Mistakes	Business & Risk & Strategy	1
57	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
58	Pete Liegl’s Performance	Business & Financial	0.475
59	GEICO Restructuring	Business & Strategy	0.46
60	GEICO Transformation	Business Strategy & Operational
61	Berkshire’s Recent Mistakes	Business & Risk & Strategy	0.536
62	Property Casualty Pricing Surge	Financial & Strategic	0.522
63	Convective Storm Damage	Financial Risk & Strategic	0.652
64	Berkshire’s Recent Challenges	Business & Risk & Strategy	0.55
65	Berkshire’s Recent Financial Challenges	Financial & Strategic
66	Insurance Losses & Strategic Risks	Business & Risk & Strategy	0.559
67	Berkshire’s Strategic Mistakes	Business & Risk & Strategy	0.747
68	Berkshire’s Financial Performance	Financial & Strategic	0.74
69	Berkshire’s Recent Mistakes	Business & Risk & Strategy
70	Berkshire’s Financial Performance	Financial & Strategic
71	Berkshire’s Recent Mistakes	Business & Risk & Strategy
72	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
73	Berkshire’s Recent Mistakes	Business & Risk & Strategy
74	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
75	Berkshire’s Recent Mistakes & Challenges	Business & Risk & Strategy	0.923
76	Berkshire’s Recent Mistakes	Business & Risk & Strategy
77	Berkshire’s Recent Mistakes & Strategic Challenges	Business & Risk & Strategy
78	Berkshire’s Recent Mistakes & Challenges	Business & Risk & Strategy
79	Berkshire’s Recent Mistakes & Strategic Challenges	Business & Risk & Strategy	0.981
80	Berkshire’s Recent Mistakes	Business & Risk & Strategy
81	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
82	Berkshire’s Recent Mistakes	Business & Risk & Strategy
83	Berkshire’s Recent Mistakes	Business & Risk & Strategy
84	Berkshire’s Recent Mistakes	Business & Risk & Strategy
85	Berkshire’s Recent Mistakes	Business & Risk & Strategy
86	Berkshire’s Mistakes	Business & Risk & Strategy	0.982
87	Berkshire’s Early Mistakes	Business & Risk & Strategy
88	Berkshire’s Early Mistakes	Business & Risk & Strategy
89	Berkshire’s Early Mistakes	Business & Risk & Strategy
90	Berkshire’s Strategic Missteps	Business & Risk & Strategy
91	Berkshire’s Tax Burden	Financial & Strategic	0.697
92	Berkshire’s Tax Burden	Financial & Risk & Strategy
93	Berkshire’s Recent Mistakes	Business & Risk & Strategy
94	Berkshire’s Recent Mistakes	Business & Risk & Strategy
95	Berkshire’s Recent Mistakes	Business & Risk & Strategy
96	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
97	Berkshire’s Recent Mistakes	Business & Risk & Strategy
98	Berkshire’s Financial Challenges	Financial & Strategic	0.742
99	Berkshire’s Recent Mistakes	Business & Risk & Strategy	0.742
100	Berkshire’s Financial Mistakes	Financial & Strategic
101	Berkshire’s Recent Mistakes	Business & Risk & Strategy
102	Berkshire’s Recent Mistakes	Business & Risk & Strategy
103	Berkshire’s Recent Mistakes	Business & Risk & Strategy
104	Berkshire’s Recent Mistakes	Business & Risk & Strategy
105	Berkshire’s Recent Mistakes & Challenges	Business & Risk & Strategy
106	Berkshire’s Recent Mistakes & Challenges	Business & Risk & Strategy
107	Berkshire’s Recent Mistakes	Business & Risk & Strategy	0.964
108	Berkshire’s Recent Mistakes	Business & Risk & Strategy
109	Berkshire’s Recent Mistakes	Business & Risk & Strategy
110	Berkshire’s Recent Mistakes	Business & Risk & Strategy
111	Berkshire’s Recent Mistakes	Business & Risk & Strategy
112	Berkshire’s Recent Mistakes	Business & Risk & Strategy
113	Berkshire’s Recent Mistakes	Business & Risk & Strategy
114	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
115	Berkshire’s Strategic Setbacks	Business & Risk & Strategy
116	Berkshire’s Recent Mistakes	Business & Risk & Strategy
117	Berkshire’s Recent Mistakes	Business & Risk & Strategy
118	Berkshire’s Recent Mistakes	Business & Risk & Strategy
119	Berkshire’s Strategic Challenges	Business & Risk & Strategy
120	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
121	Berkshire’s Recent Mistakes	Business & Risk & Strategy
122	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
123	Berkshire’s Recent Mistakes	Business & Risk & Strategy
124	Berkshire’s Recent Mistakes	Business & Risk & Strategy
125	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
126	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
127	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
128	Berkshire’s Strategic Challenges	Business & Risk & Strategy
129	Berkshire’s Recent Mistakes	Business & Risk & Strategy	0.851
130	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
131	Berkshire’s Strategic Mistakes	Business & Risk & Strategy
132	Berkshire’s Recent Mistakes	Business & Risk & Strategy
133	Berkshire’s Recent Mistakes	Business & Risk & Strategy
134	Berkshire’s Recent Mistakes	Business & Risk & Strategy
135	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
136	Berkshire’s Early Struggles	Business & Risk & Strategy
137	Berkshire’s Recent Mistakes	Business & Risk & Strategy
138	Berkshire’s Recent Mistakes	Business & Risk & Strategy
139	Berkshire’s Recent Mistakes	Business & Risk & Strategy
140	Berkshire’s Recent Mistakes	Business & Risk & Strategy
141	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
142	Berkshire’s Recent Mistakes	Business & Risk & Strategy
143	Berkshire’s Recent Mistakes & Challenges	Business & Risk & Strategy
144	Berkshire’s Early Struggles	Business & Risk & Strategy
145	Berkshire’s Recent Mistakes	Business & Risk & Strategy	0.86
146	Berkshire’s Recent Mistakes	Business & Risk & Strategy
147	Berkshire’s Recent Mistakes	Business & Risk & Strategy
148	Berkshire’s Recent Mistakes	Business & Risk & Strategy
149	Berkshire’s Recent Mistakes	Business & Risk & Strategy
150	Berkshire’s Recent Mistakes	Business & Risk & Strategy
151	Berkshire’s Recent Mistakes	Business & Risk & Strategy
152	Berkshire’s Recent Mistakes	Business & Risk & Strategy
153	Berkshire’s Recent Mistakes	Business & Risk & Strategy
154	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
155	Berkshire’s Recent Mistakes	Business & Risk & Strategy
156	Berkshire’s Recent Mistakes	Business & Risk & Strategy
157	Berkshire’s Recent Mistakes	Business & Risk & Strategy
158	Berkshire’s Recent Mistakes	Business & Risk & Strategy
159	CEOs’ Mistakes	Business Strategy & Risk
160	Berkshire’s Recent Mistakes	Business & Risk & Strategy
161	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
162	Berkshire’s Recent Mistakes	Business & Risk & Strategy
163	Berkshire’s Recent Mistakes	Business & Risk & Strategy
164	Berkshire’s Strategic Challenges	Business & Risk & Strategy	0.851
165	Berkshire’s Recent Mistakes	Business & Risk & Strategy
166	Berkshire’s Recent Mistakes	Business & Risk & Strategy
167	Berkshire’s Recent Mistakes	Business & Risk & Strategy
168	Berkshire’s Recent Mistakes	Business & Risk & Strategy
169	Berkshire’s Recent Mistakes	Business & Risk & Strategy
170	Berkshire’s Recent Mistakes	Business & Risk & Strategy
171	Berkshire’s Recent Mistakes	Business & Risk & Strategy
172	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
173	Berkshire’s Recent Mistakes	Business & Risk & Strategy
174	Berkshire’s Recent Mistakes	Business & Risk & Strategy
175	Berkshire’s Recent Mistakes	Business & Risk & Strategy
176	Berkshire’s Recent Mistakes	Business & Risk & Strategy
177	Berkshire’s Recent Mistakes	Business & Risk & Strategy
178	Berkshire’s Recent Mistakes	Business & Risk & Strategy
179	Berkshire’s Recent Mistakes	Business & Risk & Strategy
180	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
181	Berkshire’s Recent Mistakes	Business & Risk & Strategy
182	Strategic Setbacks	Business & Risk & Strategy
183	Recent Mistakes in Berkshire’s Operations	Business & Risk & Strategy	0.743
184	Hurricane, Tornado, and Wildfire Risks	Financial & Strategic	0.547
185	Strategic Mistakes	Business & Risk & Strategy	0.55
186	Auto Insurance Transition	Financial & Strategic	0.481
187	Berkshire’s Recent Mistakes	Business & Risk & Strategy	0.49
188	Berkshire’s Recent Mistakes	Business & Risk & Strategy
189	Berkshire’s Recent Mistakes	Business & Risk & Strategy
190	Berkshire’s Recent Mistakes	Business & Risk & Strategy
191	Berkshire’s Recent Mistakes	Business & Risk & Strategy
192	Berkshire’s Recent Mistakes	Business & Risk & Strategy
193	Berkshire’s Recent Mistakes	Business & Risk & Strategy
194	Berkshire’s Recent Mistakes	Business & Risk & Strategy
195	Berkshire’s Recent Mistakes	Business & Risk & Strategy
196	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
197	Berkshire’s Recent Mistakes	Business & Risk & Strategy
198	Berkshire’s Recent Mistakes	Business & Risk & Strategy
199	Berkshire’s Strategic Shifts	Business & Strategy
200	Berkshire’s Recent Mistakes	Business & Risk & Strategy
201	Berkshire’s Recent Mistakes	Business & Risk & Strategy
202	Berkshire’s Recent Mistakes	Business & Risk & Strategy
203	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
204	Berkshire’s Recent Mistakes	Business & Risk & Strategy
205	Berkshire’s Recent Mistakes	Business & Risk & Strategy
206	Berkshire’s Recent Mistakes	Business & Risk & Strategy
207	Berkshire’s Recent Mistakes	Business & Risk & Strategy
208	Berkshire’s Recent Mistakes	Business & Risk & Strategy
209	Berkshire’s Recent Mistakes	Business & Risk & Strategy
210	Berkshire’s Recent Mistakes	Business & Risk & Strategy
211	Berkshire’s Recent Mistakes	Business & Risk & Strategy	1
212	Berkshire’s Recent Mistakes	Business & Risk & Strategy
213	Berkshire’s Recent Mistakes	Business & Risk & Strategy
214	Strategic Missteps	Business & Risk & Strategy
215	Strategic Missteps	Business & Risk & Strategy
216	Strategic Missteps	Business & Risk & Strategy
217	Strategic Missteps	Business & Risk & Strategy
218	Strategic Challenges	Business & Risk & Strategy	0.803
219	Strategic Missteps	Business & Risk & Strategy
220	Strategic Setbacks	Business & Risk & Strategy
221	Recent Mistakes	Business & Risk & Strategy
222	Recent Mistakes	Business & Risk & Strategy
223	Recent Mistakes	Business & Risk & Strategy
224	Recent Mistakes	Business & Risk & Strategy
225	Recent Mistakes	Business & Risk & Strategy
226	Recent Mistakes	Business & Risk & Strategy	1
227	Recent Mistakes	Business & Risk & Strategy
228	Strategic Missteps	Business & Risk & Strategy
229	Recent Mistakes	Business & Risk & Strategy
230	Strategic Mishaps	Business & Risk & Strategy
231	Recent Mistakes	Business & Risk & Strategy
232	Recent Mistakes	Business & Risk & Strategy
233	Recent Mistakes	Business & Risk & Strategy
234	Financial Risks & Challenges	Business & Risk & Strategy	0.671
235	Recent Business Mistakes	Business & Risk & Strategy	0.694
236	Recent Mistakes	Business & Risk & Strategy
237	Strategic Challenges	Business & Risk & Strategy	0.678
238	Strategic Challenges	Business & Risk & Strategy
239	Recent Business Mistakes	Business & Risk & Strategy	0.671
240	Strategic Challenges	Business & Risk & Strategy	0.671
241	Recent Business Mistakes	Business & Risk & Strategy	0.671
242	Strategic Challenges	Business & Risk & Strategy	0.671
243	Recent Business Mistakes	Business & Risk & Strategy	0.671
244	Recent Business Mistakes	Business & Risk & Strategy
245	Recent Business Mistakes	Business & Risk & Strategy
246	Recent Business Mistakes	Business & Risk & Strategy
247	Recent Business Mistakes	Business & Risk & Strategy
248	Recent Business Mistakes	Business & Risk & Strategy
249	Recent Business Mistakes	Business & Risk & Strategy	1
250	Recent Business Mistakes	Business & Risk & Strategy
251	Recent Business Mistakes	Business & Risk & Strategy
252	Strategic Challenges	Business & Risk & Strategy	0.671
253	Recent Business Mistakes	Business & Risk & Strategy	0.671
254	Recent Business Mistakes	Business & Risk & Strategy
255	Recent Business Mistakes	Business & Risk & Strategy
256	Recent Business Mistakes	Business & Risk & Strategy
257	Recent Business Mistakes	Business & Risk & Strategy
258	Recent Business Mistakes	Business & Risk & Strategy	1

Implementation Code

import re
import json
import ollama
import logging
import numpy as np
from typing import List
from pathlib import Path
from dataclasses import dataclass
from pydantic import BaseModel, Field
from sklearn.metrics.pairwise import cosine_similarity

logging.basicConfig(level=logging.DEBUG)
log = logging.getLogger(__name__)

# Dont show logging for httpcore or httpx
logging.getLogger("httpcore").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("pdfminer").setLevel(logging.WARNING)
logging.getLogger("urllib3").setLevel(logging.WARNING)


@dataclass
class ChunkingConfig:
 max_chunk_size: int = 1000 # characters
 min_chunk_size: int = 100 # characters
 similarity_threshold: float = 0.15 # cosine distance threshold for splitting (higher means more splits)
 ollama_model: str = "nomic-embed-text:v1.5" # Ollama model for embeddings
 overlap_sentences: int = 1 # sentences to overlap between chunks


@dataclass
class SectioningConfig:
 # Model settings
 ollama_model: str = "gemma3:1b"
 max_tokens: int = 200
 temperature: float = 0.1 # Low for consistency

 # Similarity thresholds
 section_merge_threshold: float = 0.75
 title_weight: float = 0.6
 category_weight: float = 0.4

 # Size constraints
 max_section_size: int = 3000
 min_section_size: int = 300


class SimpleChunkMetadata(BaseModel):
 title: str = Field(..., description="Concise descriptive title for this content")
 category: str = Field(..., description="Broad topic category")


class Subsection(BaseModel):
 subsection_title: str
 chunks: List[str]
 combined_content: str
 chunk_count: int
 total_length: int


class DocumentSection(BaseModel):
 section_title: str
 category: str
 subsections: List[Subsection]
 total_chunks: int
 total_length: int


class SemanticChunker:
 def __init__(self, config: ChunkingConfig):
 self.config = config

 def get_embedding(self, text: str) -> np.ndarray:
 """Get embedding from Ollama using the Python client"""
 # Ensure the text is not empty, as empty strings might cause issues with embeddings
 if not text.strip():
 return np.zeros(768) # Return a zero vector for empty text, assuming 768 dimensions for nomic-embed-text
 response = ollama.embeddings(model=self.config.ollama_model, prompt=text)
 return np.array(response["embedding"])

 def split_into_sentences(self, text: str) -> List[str]:
 """Split text into sentences using regex"""
 # Improved regex to handle common sentence endings while avoiding splitting on abbreviations (e.g., Mr. Smith)
 # It looks for . ! or ? followed by whitespace and an uppercase letter, but not if preceded by a common abbreviation pattern.
 # This is a common pattern for sentence splitting.
 sentence_pattern = r"(?<!\b[A-Z]\.)(?<!\b[A-Z][a-z]\.)(?<=[.!?])\s+(?=[A-Z])"
 sentences = re.split(sentence_pattern, text)
 return [s.strip() for s in sentences if s.strip()]

 def chunk_sentences(self, sentences: List[str]) -> List[str]:
 """
 Chunks sentences based on semantic similarity, respecting max_chunk_size,
 and adding overlap.
 """
 if not sentences:
 return []

 chunks: List[str] = []
 current_chunk_sentences: List[str] = []
 current_chunk_char_length = 0

 # Get embeddings for all sentences
 sentence_embeddings = [self.get_embedding(s) for s in sentences]

 for i, sentence in enumerate(sentences):
 sentence_length = len(sentence)

 # Check if adding the current sentence would exceed max_chunk_size
 # and if we have enough content to form a valid chunk
 if current_chunk_char_length + sentence_length > self.config.max_chunk_size and current_chunk_sentences:
 # If we're at the beginning of a potential new chunk and the first sentence itself is too long,
 # we'll add it as a standalone chunk and handle it.
 if current_chunk_char_length == 0:
 chunks.append(sentence)
 current_chunk_sentences = []
 current_chunk_char_length = 0
 continue # Move to the next sentence
 else:
 # Finalize the current chunk
 chunks.append(" ".join(current_chunk_sentences))
 # Reset for the next chunk, adding overlap
 current_chunk_sentences = sentences[max(0, i - self.config.overlap_sentences) : i]
 current_chunk_char_length = sum(len(s) for s in current_chunk_sentences)

 # Add sentence to current chunk
 current_chunk_sentences.append(sentence)
 current_chunk_char_length += sentence_length

 # Semantic split check (only if there's more than one sentence in current_chunk_sentences)
 if len(current_chunk_sentences) > 1:
 # Compare the last two sentences in the current_chunk_sentences
 # We're interested in the similarity between the last added sentence and the one before it
 # to detect a potential topic shift at the boundary.
 embed1 = sentence_embeddings[i - 1] # Embedding of the sentence before the current one
 embed2 = sentence_embeddings[i] # Embedding of the current sentence

 # Calculate cosine distance (1 - cosine_similarity)
 # A higher distance means less similarity.
 if np.dot(embed1, embed2) == 0 and np.linalg.norm(embed1) == 0 and np.linalg.norm(embed2) == 0:
 # Handle case where both embeddings are zero vectors (e.g., from empty strings), distance is 0
 distance = 0
 elif np.linalg.norm(embed1) == 0 or np.linalg.norm(embed2) == 0:
 # If one is zero and other is not, they are dissimilar
 distance = 1.0
 else:
 distance = 1 - cosine_similarity(embed1.reshape(1, -1), embed2.reshape(1, -1))[0][0]

 if distance > self.config.similarity_threshold:
 log.debug(f"Semantic split detected between sentences {i - 1} and {i} with distance {distance:.4f}")
 # Semantic split detected!
 # If the chunk is long enough, finalize it before the current sentence.
 if current_chunk_char_length - sentence_length >= self.config.min_chunk_size:
 # Append chunk excluding the current sentence, which starts a new one
 chunks.append(" ".join(current_chunk_sentences[:-1]))
 # Reset for the next chunk, adding overlap
 current_chunk_sentences = sentences[max(0, i - self.config.overlap_sentences) : i + 1]
 current_chunk_char_length = sum(len(s) for s in current_chunk_sentences)
 # If the chunk is too short, we don't split yet and continue to build it.
 # This prevents very small chunks due to minor semantic shifts.

 # Add any remaining sentences as the last chunk
 if current_chunk_sentences:
 chunks.append(" ".join(current_chunk_sentences))

 # Post-processing: Ensure no chunks are empty and optionally merge very small chunks
 final_chunks = [chunk for chunk in chunks if chunk.strip()]
 return final_chunks


class SectionGrouper:
 def __init__(self, sectioning_config: SectioningConfig, embedding_model: str = "nomic-embed-text:v1.5"):
 self.config = sectioning_config
 self.embedding_model = embedding_model

 def generate_chunk_metadata(self, chunk: str, previous_metadata) -> SimpleChunkMetadata:
 """Generate structured metadata for a chunk using Ollama"""

 # Prepare previous metadata as a string for context
 metadata = ""
 if len(previous_metadata) > 0:
 metadata = (
 "\n".join(
 f"{i + 1}. Title: {meta.title}, Category: {meta.category}"
 for i, (chunk, meta) in enumerate(previous_metadata)
 )
 if previous_metadata
 else ""
 )

 if metadata != "":
 metadata = f"Previous metadata:\n{metadata}\n\n"

 prompt = f"""Analyze the following text and provide structured metadata:

Text: {chunk}

{metadata}

Generate a JSON response with:
- title: A concise 3-8 word descriptive title
- category: A broad 1-3 word topic category

Be specific and descriptive but concise. Keep consistency in section titles and categories across similar content."""

 try:
 response = ollama.generate(
 model=self.config.ollama_model,
 prompt=prompt,
 format=SimpleChunkMetadata.model_json_schema(),
 options={"temperature": self.config.temperature, "num_predict": self.config.max_tokens},
 keep_alive=True,
 )

 # Parse the JSON response
 metadata_dict = json.loads(response["response"])
 return SimpleChunkMetadata(**metadata_dict)

 except Exception as e:
 log.warning(f"Failed to generate metadata for chunk, using fallback: {e}")
 # Fallback metadata
 return SimpleChunkMetadata(
 title=f"Section {hash(chunk[:100]) % 1000}",
 category="General",
 )

 def get_embedding(self, text: str) -> np.ndarray:
 """Get embedding for text"""
 if not text.strip():
 return np.zeros(768)
 response = ollama.embeddings(model=self.embedding_model, prompt=text)
 return np.array(response["embedding"])

 def calculate_similarity(self, metadata1: SimpleChunkMetadata, metadata2: SimpleChunkMetadata) -> float:
 """Calculate weighted similarity between two chunk metadata objects"""

 # Get embeddings for each field
 title1_emb = self.get_embedding(metadata1.title)
 title2_emb = self.get_embedding(metadata2.title)

 category1_emb = self.get_embedding(metadata1.category)
 category2_emb = self.get_embedding(metadata2.category)

 # Calculate cosine similarities
 def safe_cosine_similarity(emb1, emb2):
 if np.linalg.norm(emb1) == 0 or np.linalg.norm(emb2) == 0:
 return 0.0
 return cosine_similarity(emb1.reshape(1, -1), emb2.reshape(1, -1))[0][0]

 title_sim = safe_cosine_similarity(title1_emb, title2_emb)
 category_sim = safe_cosine_similarity(category1_emb, category2_emb)

 # Weighted similarity
 weighted_sim = self.config.title_weight * title_sim + self.config.category_weight * category_sim

 return weighted_sim

 def group_chunks_into_sections(self, chunks: List[str]) -> List[DocumentSection]:
 """Group chunks into coherent sections based on metadata similarity"""

 if not chunks:
 return []

 log.info(f"Generating metadata for {len(chunks)} chunks...")

 # Stage 1: Generate metadata for all chunks
 chunk_metadata = []
 for i, chunk in enumerate(chunks):
 # We pass in the last 5 metadata items as context for the next chunk
 metadata = self.generate_chunk_metadata(
 chunk, chunk_metadata[:-5] if len(chunk_metadata) >= 5 else chunk_metadata
 )
 chunk_metadata.append((chunk, metadata))
 log.debug(f"Chunk {i + 1}: {metadata.title} | {metadata.category}")

 log.info("Grouping chunks into sections...")

 # Stage 2: Group chunks based on similarity
 sections = []
 current_section_chunks = [(chunk_metadata[0][0], chunk_metadata[0][1])]
 current_section_category = chunk_metadata[0][1].category

 for i in range(1, len(chunk_metadata)):
 chunk, metadata = chunk_metadata[i]
 prev_metadata = chunk_metadata[i - 1][1]

 # Calculate similarity with previous chunk
 similarity = self.calculate_similarity(metadata, prev_metadata)

 # Check if we should start a new section
 should_split = (
 # Similarity is below threshold so the current chunk is likely different
 # from the previous section
 similarity < self.config.section_merge_threshold
 # Or if the current section is too long
 or sum(len(c[0]) for c in current_section_chunks) + len(chunk) > self.config.max_section_size
 )

 if should_split and len(current_section_chunks) > 0:
 # Finalize current section
 section = self._create_section(current_section_chunks, current_section_category)
 sections.append(section)

 # Start new section
 current_section_chunks = [(chunk, metadata)]
 current_section_category = metadata.category

 log.debug(f"New section started at chunk {i + 1}, similarity: {similarity:.3f}")
 else:
 # Add to current section
 current_section_chunks.append((chunk, metadata))

 # Add final section
 if current_section_chunks:
 section = self._create_section(current_section_chunks, current_section_category)
 sections.append(section)

 log.info(f"Created {len(sections)} sections from {len(chunks)} chunks")
 return sections

 def _create_section(self, section_chunks: List[tuple], category: str) -> DocumentSection:
 """Create a DocumentSection from a list of (chunk, metadata) tuples"""

 # Group chunks into subsections based on title similarity
 subsections = []
 current_subsection = []
 current_title = section_chunks[0][1].title

 for chunk, metadata in section_chunks:
 if not current_subsection:
 current_subsection = [(chunk, metadata)]
 current_title = metadata.title
 else:
 # Check if we should group with current subsection
 title_sim = self.calculate_similarity(metadata, current_subsection[-1][1])

 if title_sim > 0.8: # High threshold for subsection grouping
 current_subsection.append((chunk, metadata))
 else:
 # Create subsection from current group
 subsection = self._create_subsection(current_subsection, current_title)
 subsections.append(subsection)

 # Start new subsection
 current_subsection = [(chunk, metadata)]
 current_title = metadata.title

 # Add final subsection
 if current_subsection:
 subsection = self._create_subsection(current_subsection, current_title)
 subsections.append(subsection)

 # Create section title from most common category and representative title
 section_title = self._generate_section_title(section_chunks)

 return DocumentSection(
 section_title=section_title,
 category=category,
 subsections=subsections,
 total_chunks=len(section_chunks),
 total_length=sum(len(chunk) for chunk, _ in section_chunks),
 )

 def _create_subsection(self, subsection_chunks: List[tuple], title: str) -> Subsection:
 """Create a Subsection from a list of (chunk, metadata) tuples"""
 chunks = [chunk for chunk, _ in subsection_chunks]
 combined_content = "\n\n".join(chunks)

 return Subsection(
 subsection_title=title,
 chunks=chunks,
 combined_content=combined_content,
 chunk_count=len(chunks),
 total_length=len(combined_content),
 )

 def _generate_section_title(self, section_chunks: List[tuple]) -> str:
 """Generate a representative title for the section"""
 # Use the title from the first chunk, or create from category
 if section_chunks:
 first_metadata = section_chunks[0][1]
 if len(section_chunks) == 1:
 return first_metadata.title
 else:
 # For multi-chunk sections, use category-based title
 return f"{first_metadata.category} Overview"
 return "Untitled Section"


def pdf_to_markdown(pdf_file: Path) -> str:
 from markitdown import MarkItDown

 md = MarkItDown()
 result = md.convert(pdf_file)
 text_content = result.text_content
 # Keep only ascii characters
 text_content = "".join(c for c in text_content if ord(c) < 128)
 # Remove leading and trailing whitespace
 text_content = text_content.strip()
 # Remove multiple newlines
 text_content = "\n".join(line.strip() for line in text_content.splitlines() if line.strip())
 return text_content


def get_pdf():
 import requests

 # Sample text chunks
 url = "https://www.berkshirehathaway.com/letters/2024ltr.pdf"
 response = requests.get(url)
 if response.status_code != 200:
 raise ValueError(f"Failed to download PDF: {response.status_code}")
 pdf_file = Path("./example_cache/2024ltr.pdf")
 pdf_file.parent.mkdir(parents=True, exist_ok=True)
 pdf_file.write_bytes(response.content)
 # Convert PDF to markdown text
 text = pdf_to_markdown(pdf_file)
 return text


# Usage example
def main():
 # Example usage
 sample_text = get_pdf()

 # Initialize chunker
 chunking_config = ChunkingConfig(
 max_chunk_size=750,
 min_chunk_size=150,
 similarity_threshold=0.25,
 ollama_model="nomic-embed-text:v1.5",
 overlap_sentences=2,
 )

 # Initialize section grouper
 sectioning_config = SectioningConfig(
 ollama_model="gemma3:1b", section_merge_threshold=0.75, max_section_size=3000, min_section_size=300
 )

 # Stage 1: Semantic chunking
 chunker = SemanticChunker(chunking_config)
 sentences = chunker.split_into_sentences(sample_text)
 chunks = chunker.chunk_sentences(sentences)

 print(f"Stage 1 Complete - Generated {len(chunks)} semantic chunks")

 # Stage 2: Section grouping
 grouper = SectionGrouper(sectioning_config)
 sections = grouper.group_chunks_into_sections(chunks)

 # Write sections to JSON file
 output_file = Path("./example_cache/sections.json")
 output_file.parent.mkdir(parents=True, exist_ok=True)
 with output_file.open("w", encoding="utf-8") as f:
 json.dump([section.model_dump() for section in sections], f, indent=2, ensure_ascii=False)
 log.info(f"Sections written to {output_file}")

 print(f"Stage 2 Complete - Generated {len(sections)} sections")

 # Display results
 print("\n" + "=" * 80)
 print("DOCUMENT STRUCTURE")
 print("=" * 80)

 for i, section in enumerate(sections, 1):
 print(f"\nSECTION {i}: {section.section_title}")
 print(f"Category: {section.category}")
 print(f"Total chunks: {section.total_chunks} | Length: {section.total_length} chars")
 print("-" * 60)

 for j, subsection in enumerate(section.subsections, 1):
 print(f" {i}.{j} {subsection.subsection_title}")
 print(f" Chunks: {subsection.chunk_count} | Length: {subsection.total_length} chars")

 # Show first chunk preview
 if subsection.chunks:
 preview = (
 subsection.chunks[0][:200] + "..." if len(subsection.chunks[0]) > 200 else subsection.chunks[0]
 )
 print(f" Preview: {preview}")
 print()

Python Dependency Injection

Sun, 15 Jun 2025 08:14:56 -0400

Dependency Injection (DI) is a software design pattern that lets you pass instances of a service rather than creating them directly within a class or function. The FastAPI framework provides a neat way to pass dependencies using Python’s type hints and the Depends function. The fast-depends package extracts the FastAPI code and strips out all the web framework-specific code into a small library that can be used in any Python project.

Here is an example of how to use dependency injection using the custom single-file implementation of dependency injection at the bottom of this page.

import typing as t

def get_settings() -> Settings:
 return Settings()

@inject
def get_db(
 settings: t.Annotated[Settings, Depends(get_settings)],
) -> DatabaseConnection:
 db = DatabaseConnection(settings)
 return db

@inject
def compute_something(
 db: t.Annotated[DatabaseConnection, Depends(get_db)]
):
 # Use db connection here
 pass

# Call the function with dependencies injected
result = compute_something() # Automatically resolves dependencies

db = get_db() # You can also call the dependency directly
result = compute_something(db=db) # Pass the db directly if needed

The get_* functions return instances of the required services, and the Depends function is used to declare dependencies. The inject decorator enables dependency injection for the function by automatically resolving the dependencies when the function is called. The default behavior is to cache the results of dependencies, so if a dependency is called multiple times, it will return the cached result instead of executing the function again (e.g. Settings and DB connection are created once and reused).

This pattern allows for better separation of concerns (instantiatiating dependencies outside of the function) and makes unit testing easier. For example, we can create different implementations of the database and settings classes and pass them to the compute_something function without changing its signature.

Below is the full implementation of dependency injection as a standalone Python module. You can copy it into a depends.py file and use it in your projects.

Key Components:

Depends: Marks a parameter as a dependency to be injected
inject: Decorator that enables dependency injection for a sync or async function
CustomField: Base class for creating custom parameter extractors
dependency_provider: Global provider for managing dependency overrides

Features:

Automatic dependency resolution and injection
Support for both sync and async functions
Dependency caching (can be disabled per dependency)
Type validation and casting using Pydantic
Context manager support for resource management
Custom field extractors for complex parameter handling
Dependency override system for testing and configuration

Table of Contents:

Usage
Implementation Code

Usage

Simple Dependency Injection

Dependencies can be injected into functions using the Depends class and the inject decorator.

def get_database():
 return "database_connection"

def get_user(db: str = Depends(get_database)):
 return f"user_from_{db}"

@inject
def handler(user: str = Depends(get_user)):
 return f"Hello, {user}!"

result = handler() # "Hello, user_from_database_connection!"
print(result)

Async Dependencies

async def get_async_db():
 await asyncio.sleep(0.1)
 return "async_database"

@inject
async def async_handler(db: str = Depends(get_async_db)):
 return f"DB: {db}"

result = await async_handler() # "DB: async_database"
print(result)

Dependency Caching

Caching is enabled by default, meaning that if a dependency is called multiple times within the same request, it will return the cached result instead of executing the function again.

call_count = 0

def expensive_operation():
 global call_count
 call_count += 1
 return f"result_{call_count}"

@inject
def handler(
 a: str = Depends(expensive_operation),
 b: str = Depends(expensive_operation), # Same result due to caching
):
 return f"{a}, {b}"

result = handler() # "result_1, result_1"
print(result) # Output: "result_1, result_1"

Disable Caching

The use_cache parameter can be set to False to disable caching for specific dependencies.

@inject
def handler(
 a: str = Depends(expensive_operation, use_cache=False),
 b: str = Depends(expensive_operation, use_cache=False),
):
 return f"{a}, {b}"

result = handler() # "result_1, result_2"
print(result) # Output: "result_1, result_2"

Custom Fields

class HeaderExtractor(CustomField):
 def __init__(self, header_name: str):
 super().__init__()
 self.header_name = header_name

 def use(self, **kwargs: t.Any) -> t.Dict[str, t.Any]:
 # Extract from some global context
 kwargs[self.param_name] = f"header_value_for_{self.header_name}"
 return kwargs

@inject
def api_handler(
 auth: str = HeaderExtractor("Authorization"),
 content_type: str = HeaderExtractor("Content-Type"),
):
 return {"auth": auth, "content_type": content_type}

Dependency Overrides

def original_dep():
 return "original"

def override_dep():
 return "overridden"

@inject
def handler(value: str = Depends(original_dep)):
 return value

# Override dependency
dependency_provider.override(original_dep, override_dep)
result = handler() # "overridden"

# Clear overrides
dependency_provider.clear()
result = handler() # "original"

Generator Dependencies (Context Managers)

A dependency can be a generator function, which allows for resource management (like opening and closing database connections). In this example, the database_session function is a context manager that opens a database connection and closes it after use. The Depends decorator will handle the context management automatically.

def database_session() -> t.Generator[str, None, None]:
 print("Opening connection")
 yield "db_session"
 print("Closing connection")

@inject
def handler(db: str = Depends(database_session)):
 return f"Using {db}"

result = handler()
print(result) # "Using db_session"
# Output: Opening connection
# Output: Closing connection
# Returns: "Using db_session"

Type Validation with Pydantic

Arguments can be annotated with Pydantic models for automatic validation and casting. In this example, the get_user_id function returns a string, but it will be cast to an integer when injected. The Depends decorator will handle the type casting automatically if cast=True is set (default behavior). The Annotated type from typing_extensions is used to specify the type and the dependency.

from typing import Annotated

def get_user_id() -> int:
 return "123" # Wrong type!

@inject
def handler(user_id: Annotated[int, Depends(get_user_id)]):
 return f"User ID: {user_id}"

result = handler() # user_id will be cast to int(123)

Disable Type Casting

If you want to disable type casting for a specific dependency, you can set cast=False in the inject decorator.

@inject(cast=False)
def handler(user_id: Annotated[int, Depends(get_user_id)]):
 return f"User ID type: {type(user_id)}"

result = handler() # user_id remains as string "123"

Implementation Code

Requirements:

Python 3.11+
anyio for async I/O operations
pydantic for data validation and settings management
typing_extensions for type annotations

import anyio
import asyncio
import inspect
import functools
import typing as t
from abc import ABC
from copy import deepcopy
from itertools import chain
from pydantic import ConfigDict
from collections import namedtuple
from functools import wraps, partial
from pydantic import BaseModel, create_model
from typing_extensions import Annotated, ParamSpec, get_args, get_origin
from pydantic._internal._typing_extra import try_eval_type as evaluate_forwardref
from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager

P = ParamSpec("P")
T = t.TypeVar("T")
Cls = t.TypeVar("Cls", bound="CustomField")

default_pydantic_config = {"arbitrary_types_allowed": True}


def get_config_base(config_data: t.Optional[ConfigDict] = None) -> ConfigDict:
 return config_data or ConfigDict(**default_pydantic_config)


def get_aliases(model: t.Type[BaseModel]) -> t.Tuple[str, ...]:
 return tuple(f.alias or name for name, f in model.model_fields.items())


class Depends:
 """Mark a parameter as a dependency to be injected."""

 use_cache: bool
 cast: bool

 def __init__(
 self,
 dependency: t.Callable[..., t.Any],
 *,
 use_cache: bool = True,
 cast: bool = True,
 ) -> None:
 self.dependency = dependency
 self.use_cache = use_cache
 self.cast = cast

 def __repr__(self) -> str:
 attr = getattr(self.dependency, "__name__", type(self.dependency).__name__)
 cache = "" if self.use_cache else ", use_cache=False"
 return f"{self.__class__.__name__}({attr}{cache})"


class CustomField(ABC):
 """Base class for custom field extractors."""

 param_name: t.Optional[str]
 cast: bool
 required: bool

 __slots__ = (
 "cast",
 "param_name",
 "required",
 "field",
 )

 def __init__(
 self,
 *,
 cast: bool = True,
 required: bool = True,
 ) -> None:
 self.cast = cast
 self.param_name = None
 self.required = required
 self.field = False

 def set_param_name(self: Cls, name: str) -> Cls:
 self.param_name = name
 return self

 def use(self, /, **kwargs: t.Any) -> t.Dict[str, t.Any]:
 assert self.param_name, "You should specify `param_name` before using"
 return kwargs

 def use_field(self, kwargs: t.Dict[str, t.Any]) -> None:
 raise NotImplementedError("You should implement `use_field` method.")


# Provider for dependency overrides


class Provider:
 """Provider for dependency overrides."""

 dependency_overrides: t.Dict[t.Callable[..., t.Any], t.Callable[..., t.Any]]

 def __init__(self) -> None:
 self.dependency_overrides = {}

 def clear(self) -> None:
 self.dependency_overrides = {}

 def override(
 self,
 original: t.Callable[..., t.Any],
 override: t.Callable[..., t.Any],
 ) -> None:
 self.dependency_overrides[original] = override

 @contextmanager
 def scope(
 self,
 original: t.Callable[..., t.Any],
 override: t.Callable[..., t.Any],
 ) -> t.Iterator[None]:
 self.dependency_overrides[original] = override
 yield
 self.dependency_overrides.pop(original, None)


dependency_provider = Provider()


def is_coroutine_callable(call: t.Callable[..., t.Any]) -> bool:
 if inspect.isclass(call):
 return False
 if asyncio.iscoroutinefunction(call):
 return True
 dunder_call = getattr(call, "__call__", None)
 return asyncio.iscoroutinefunction(dunder_call)


def is_gen_callable(call: t.Callable[..., t.Any]) -> bool:
 if inspect.isgeneratorfunction(call):
 return True
 dunder_call = getattr(call, "__call__", None)
 return inspect.isgeneratorfunction(dunder_call)


def is_async_gen_callable(call: t.Callable[..., t.Any]) -> bool:
 if inspect.isasyncgenfunction(call):
 return True
 dunder_call = getattr(call, "__call__", None)
 return inspect.isasyncgenfunction(dunder_call)


async def run_async(
 func: t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]],
 *args: P.args,
 **kwargs: P.kwargs,
) -> T:
 if is_coroutine_callable(func):
 return await t.cast(t.Callable[P, t.Awaitable[T]], func)(*args, **kwargs)
 else:
 return await run_in_threadpool(t.cast(t.Callable[P, T], func), *args, **kwargs)


async def run_in_threadpool(func: t.Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
 if kwargs:
 func = functools.partial(func, **kwargs)
 return await anyio.to_thread.run_sync(func, *args)


def get_typed_annotation(
 annotation: t.Any,
 globalns: t.Dict[str, t.Any],
 locals: t.Dict[str, t.Any],
) -> t.Any:
 if isinstance(annotation, str):
 annotation = t.ForwardRef(annotation)

 if isinstance(annotation, t.ForwardRef):
 annotation = evaluate_forwardref(annotation, globalns, locals)

 if get_origin(annotation) is Annotated and (args := get_args(annotation)):
 solved_args = [get_typed_annotation(x, globalns, locals) for x in args]
 annotation.__origin__, annotation.__metadata__ = solved_args[0], tuple(solved_args[1:])

 return annotation


def collect_outer_stack_locals() -> t.Dict[str, t.Any]:
 """
 Collect local variables from outer stack frames to resolve type annotations.

 This function walks up the call stack and collects all local variables
 from frames outside of this module. This is necessary for resolving
 forward references and string annotations that might reference variables
 defined in the calling code.
 """
 frame = inspect.currentframe()
 frames: t.List[t.Any] = []
 current_filename = __file__ if "__file__" in globals() else None

 while frame is not None:
 frame_filename = frame.f_code.co_filename
 # Skip frames from this module to avoid internal variables
 if current_filename is None or frame_filename != current_filename:
 frames.append(frame)
 frame = frame.f_back

 locals = {}
 for f in frames[::-1]:
 locals.update(f.f_locals)

 return locals


def get_typed_signature(call: t.Callable[..., t.Any]) -> t.Tuple[inspect.Signature, t.Any]:
 signature = inspect.signature(call)
 locals = collect_outer_stack_locals()
 call = inspect.unwrap(call)
 globalns = getattr(call, "__globals__", {})

 typed_params = [
 inspect.Parameter(
 name=param.name,
 kind=param.kind,
 default=param.default,
 annotation=get_typed_annotation(
 param.annotation,
 globalns,
 locals,
 ),
 )
 for param in signature.parameters.values()
 ]

 return inspect.Signature(typed_params), get_typed_annotation(
 signature.return_annotation,
 globalns,
 locals,
 )


async def solve_generator_async(
 *sub_args: t.Any, call: t.Callable[..., t.Any], stack: AsyncExitStack, **sub_values: t.Any
) -> t.Any:
 if is_gen_callable(call):
 cm = contextmanager_in_threadpool(contextmanager(call)(**sub_values))
 elif is_async_gen_callable(call):
 cm = asynccontextmanager(call)(*sub_args, **sub_values)
 return await stack.enter_async_context(cm)


def solve_generator_sync(
 *sub_args: t.Any, call: t.Callable[..., t.Any], stack: ExitStack, **sub_values: t.Any
) -> t.Any:
 cm = contextmanager(call)(*sub_args, **sub_values)
 return stack.enter_context(cm)


@asynccontextmanager
async def contextmanager_in_threadpool(
 cm: t.ContextManager[T],
) -> t.AsyncGenerator[T, None]:
 exit_limiter = anyio.CapacityLimiter(1)
 try:
 yield await run_in_threadpool(cm.__enter__)
 except Exception as e:
 ok = bool(await anyio.to_thread.run_sync(cm.__exit__, type(e), e, None, limiter=exit_limiter))
 if not ok:
 raise e
 else:
 await anyio.to_thread.run_sync(cm.__exit__, None, None, None, limiter=exit_limiter)


async def async_map(func: t.Callable[..., T], async_iterable: t.AsyncIterable[t.Any]) -> t.AsyncIterable[T]:
 async for i in async_iterable:
 yield func(i)


class solve_wrapper(partial[T]):
 call: t.Callable[..., T]

 def __new__(
 cls,
 func: t.Callable[..., T],
 *args: t.Any,
 **kwargs: t.Any,
 ) -> "solve_wrapper[T]":
 assert len(args) > 0, "Model should be passed as first argument"
 model = args[0]
 self = super().__new__(cls, func, *args, **kwargs)
 self.call = model.call
 return self


# Core Models

PriorityPair = namedtuple("PriorityPair", ("call", "dependencies_number", "dependencies_names"))


class ResponseModel(BaseModel, t.Generic[T]):
 response: T


class CallModel(t.Generic[P, T]):
 """Model representing a callable with dependency injection."""

 call: t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]]
 is_async: bool
 is_generator: bool
 model: t.Optional[t.Type[BaseModel]]
 response_model: t.Optional[t.Type[ResponseModel[T]]]

 params: t.Dict[str, t.Tuple[t.Any, t.Any]]
 alias_arguments: t.Tuple[str, ...]

 dependencies: t.Dict[str, "CallModel[..., t.Any]"]
 extra_dependencies: t.Iterable["CallModel[..., t.Any]"]
 sorted_dependencies: t.Tuple[t.Tuple["CallModel[..., t.Any]", int], ...]
 custom_fields: t.Dict[str, CustomField]
 keyword_args: t.Tuple[str, ...]
 positional_args: t.Tuple[str, ...]
 var_positional_arg: t.Optional[str]
 var_keyword_arg: t.Optional[str]

 use_cache: bool
 cast: bool

 __slots__ = (
 "call",
 "is_async",
 "is_generator",
 "model",
 "response_model",
 "params",
 "alias_arguments",
 "keyword_args",
 "positional_args",
 "var_positional_arg",
 "var_keyword_arg",
 "dependencies",
 "extra_dependencies",
 "sorted_dependencies",
 "custom_fields",
 "use_cache",
 "cast",
 )

 @property
 def call_name(self) -> str:
 call = inspect.unwrap(self.call)
 return getattr(call, "__name__", type(call).__name__)

 @property
 def flat_params(self) -> t.Dict[str, t.Tuple[t.Any, t.Any]]:
 params = self.params
 for d in (*self.dependencies.values(), *self.extra_dependencies):
 params.update(d.flat_params)
 return params

 @property
 def flat_dependencies(
 self,
 ) -> t.Dict[
 t.Callable[..., t.Any],
 t.Tuple["CallModel[..., t.Any]", t.Tuple[t.Callable[..., t.Any], ...]],
 ]:
 flat: t.Dict[
 t.Callable[..., t.Any],
 t.Tuple[CallModel[..., t.Any], t.Tuple[t.Callable[..., t.Any], ...]],
 ] = {}

 for i in (*self.dependencies.values(), *self.extra_dependencies):
 flat.update(
 {
 i.call: (
 i,
 tuple(j.call for j in i.dependencies.values()),
 )
 }
 )
 flat.update(i.flat_dependencies)

 return flat

 def __init__(
 self,
 /,
 call: t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]],
 model: t.Optional[t.Type[BaseModel]],
 params: t.Dict[str, t.Tuple[t.Any, t.Any]],
 response_model: t.Optional[t.Type[ResponseModel[T]]] = None,
 use_cache: bool = True,
 cast: bool = True,
 is_async: bool = False,
 is_generator: bool = False,
 dependencies: t.Optional[t.Dict[str, "CallModel[..., t.Any]"]] = None,
 extra_dependencies: t.Optional[t.Iterable["CallModel[..., t.Any]"]] = None,
 keyword_args: t.Optional[t.List[str]] = None,
 positional_args: t.Optional[t.List[str]] = None,
 var_positional_arg: t.Optional[str] = None,
 var_keyword_arg: t.Optional[str] = None,
 custom_fields: t.Optional[t.Dict[str, CustomField]] = None,
 ):
 self.call = call
 self.model = model

 if model:
 self.alias_arguments = get_aliases(model)
 else:
 self.alias_arguments = ()

 self.keyword_args = tuple(keyword_args or ())
 self.positional_args = tuple(positional_args or ())
 self.var_positional_arg = var_positional_arg
 self.var_keyword_arg = var_keyword_arg
 self.response_model = response_model
 self.use_cache = use_cache
 self.cast = cast
 self.is_async = is_async or is_coroutine_callable(call) or is_async_gen_callable(call)
 self.is_generator = is_generator or is_gen_callable(call) or is_async_gen_callable(call)

 self.dependencies = dependencies or {}
 self.extra_dependencies = extra_dependencies or ()
 self.custom_fields = custom_fields or {}

 sorted_dep: t.List[CallModel[..., t.Any]] = []
 flat = self.flat_dependencies
 for calls in flat.values():
 _sort_dep(sorted_dep, calls, flat)

 self.sorted_dependencies = tuple((i, len(i.sorted_dependencies)) for i in sorted_dep if i.use_cache)
 for name in chain(self.dependencies.keys(), self.custom_fields.keys()):
 params.pop(name, None)
 self.params = params

 def _solve(
 self,
 /,
 *args: t.Tuple[t.Any, ...],
 cache_dependencies: t.Dict[
 t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]],
 T,
 ],
 dependency_overrides: t.Optional[
 t.Dict[
 t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]],
 t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]],
 ]
 ] = None,
 **kwargs: t.Dict[str, t.Any],
 ) -> t.Generator[
 t.Tuple[t.Sequence[t.Any], t.Dict[str, t.Any], t.Callable[..., t.Any]],
 t.Any,
 T,
 ]:
 if dependency_overrides:
 call = dependency_overrides.get(self.call, self.call)
 assert self.is_async or not is_coroutine_callable(call), (
 f"You cannot use async dependency `{self.call_name}` at sync main"
 )
 else:
 call = self.call

 if self.use_cache and call in cache_dependencies:
 return cache_dependencies[call]

 kw: t.Dict[str, t.Any] = {}

 for arg in self.keyword_args:
 if (v := kwargs.pop(arg, inspect.Parameter.empty)) is not inspect.Parameter.empty:
 kw[arg] = v

 if self.var_keyword_arg is not None:
 kw[self.var_keyword_arg] = kwargs
 else:
 kw.update(kwargs)

 for arg in self.positional_args:
 if args:
 kw[arg], args = args[0], args[1:]
 else:
 break

 keyword_args: t.Iterable[str]
 if self.var_positional_arg is not None:
 kw[self.var_positional_arg] = args
 keyword_args = self.keyword_args
 else:
 keyword_args = self.keyword_args + self.positional_args
 for arg in keyword_args:
 if not self.cast and arg in self.params:
 kw[arg] = self.params[arg][1]

 if not args:
 break

 if arg not in self.dependencies:
 kw[arg], args = args[0], args[1:]

 solved_kw: t.Dict[str, t.Any]
 solved_kw = yield args, kw, call

 args_: t.Sequence[t.Any]
 if self.cast:
 assert self.model, "Cast should be used only with model"
 casted_model = self.model(**solved_kw)

 kwargs_ = {arg: getattr(casted_model, arg, solved_kw.get(arg)) for arg in keyword_args}
 if self.var_keyword_arg:
 kwargs_.update(getattr(casted_model, self.var_keyword_arg, {}))

 if self.var_positional_arg is not None:
 args_ = [getattr(casted_model, arg, solved_kw.get(arg)) for arg in self.positional_args]
 args_.extend(getattr(casted_model, self.var_positional_arg, ()))
 else:
 args_ = ()
 else:
 kwargs_ = {arg: solved_kw.get(arg) for arg in keyword_args}

 if self.var_positional_arg is not None:
 args_ = tuple(map(solved_kw.get, self.positional_args))
 else:
 args_ = ()

 response: T
 response = yield args_, kwargs_, call

 if self.cast and not self.is_generator:
 response = self._cast_response(response)

 if self.use_cache:
 cache_dependencies[call] = response

 return response

 def _cast_response(self, /, value: t.Any) -> t.Any:
 if self.response_model is not None:
 return self.response_model(response=value).response
 else:
 return value

 def solve(
 self,
 /,
 *args: t.Tuple[t.Any, ...],
 stack: ExitStack,
 cache_dependencies: t.Dict[
 t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]],
 T,
 ],
 dependency_overrides: t.Optional[
 t.Dict[
 t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]],
 t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]],
 ]
 ] = None,
 nested: bool = False,
 **kwargs: t.Dict[str, t.Any],
 ) -> T:
 cast_gen = self._solve(
 *args,
 cache_dependencies=cache_dependencies,
 dependency_overrides=dependency_overrides,
 **kwargs,
 )
 try:
 args, kwargs, _ = next(cast_gen)
 except StopIteration as e:
 cached_value: T = e.value
 return cached_value

 # Heat cache and solve extra dependencies
 for dep, _ in self.sorted_dependencies:
 dep.solve(
 *args,
 stack=stack,
 cache_dependencies=cache_dependencies,
 dependency_overrides=dependency_overrides,
 nested=True,
 **kwargs,
 )

 # Always get from cache
 for dep in self.extra_dependencies:
 dep.solve(
 *args,
 stack=stack,
 cache_dependencies=cache_dependencies,
 dependency_overrides=dependency_overrides,
 nested=True,
 **kwargs,
 )

 for dep_arg, dep in self.dependencies.items():
 kwargs[dep_arg] = dep.solve(
 stack=stack,
 cache_dependencies=cache_dependencies,
 dependency_overrides=dependency_overrides,
 nested=True,
 **kwargs,
 )

 for custom in self.custom_fields.values():
 if custom.field:
 custom.use_field(kwargs)
 else:
 kwargs = custom.use(**kwargs)

 final_args, final_kwargs, call = cast_gen.send(kwargs)

 if self.is_generator and nested:
 response = solve_generator_sync(
 *final_args,
 call=call,
 stack=stack,
 **final_kwargs,
 )
 else:
 response = call(*final_args, **final_kwargs)

 try:
 cast_gen.send(response)
 except StopIteration as e:
 value: T = e.value

 if not self.cast or nested or not self.is_generator:
 return value
 else:
 return map(self._cast_response, value)

 raise AssertionError("unreachable")

 async def asolve(
 self,
 /,
 *args: t.Tuple[t.Any, ...],
 stack: AsyncExitStack,
 cache_dependencies: t.Dict[
 t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]],
 T,
 ],
 dependency_overrides: t.Optional[
 t.Dict[
 t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]],
 t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]],
 ]
 ] = None,
 nested: bool = False,
 **kwargs: t.Dict[str, t.Any],
 ) -> T:
 cast_gen = self._solve(
 *args,
 cache_dependencies=cache_dependencies,
 dependency_overrides=dependency_overrides,
 **kwargs,
 )
 try:
 args, kwargs, _ = next(cast_gen)
 except StopIteration as e:
 cached_value: T = e.value
 return cached_value

 # Heat cache and solve extra dependencies
 dep_to_solve: t.List[t.Callable[..., t.Awaitable[t.Any]]] = []
 try:
 async with anyio.create_task_group() as tg:
 for dep, subdep in self.sorted_dependencies:
 solve = partial(
 dep.asolve,
 *args,
 stack=stack,
 cache_dependencies=cache_dependencies,
 dependency_overrides=dependency_overrides,
 nested=True,
 **kwargs,
 )
 if not subdep:
 tg.start_soon(solve)
 else:
 dep_to_solve.append(solve)
 except Exception as e:
 raise e

 for i in dep_to_solve:
 await i()

 # Always get from cache
 for dep in self.extra_dependencies:
 await dep.asolve(
 *args,
 stack=stack,
 cache_dependencies=cache_dependencies,
 dependency_overrides=dependency_overrides,
 nested=True,
 **kwargs,
 )

 for dep_arg, dep in self.dependencies.items():
 kwargs[dep_arg] = await dep.asolve(
 stack=stack,
 cache_dependencies=cache_dependencies,
 dependency_overrides=dependency_overrides,
 nested=True,
 **kwargs,
 )

 custom_to_solve: t.List[CustomField] = []

 try:
 async with anyio.create_task_group() as tg:
 for custom in self.custom_fields.values():
 if custom.field:
 tg.start_soon(run_async, custom.use_field, kwargs)
 else:
 custom_to_solve.append(custom)
 except Exception as e:
 raise e

 for j in custom_to_solve:
 kwargs = await run_async(j.use, **kwargs)

 final_args, final_kwargs, call = cast_gen.send(kwargs)

 if self.is_generator and nested:
 response = await solve_generator_async(
 *final_args,
 call=call,
 stack=stack,
 **final_kwargs,
 )
 else:
 response = await run_async(call, *final_args, **final_kwargs)

 try:
 cast_gen.send(response)
 except StopIteration as e:
 value: T = e.value

 if not self.cast or nested or not self.is_generator:
 return value
 else:
 return async_map(self._cast_response, value)

 raise AssertionError("unreachable")


def _sort_dep(
 collector: t.List["CallModel[..., t.Any]"],
 items: t.Tuple[
 "CallModel[..., t.Any]",
 t.Tuple[t.Callable[..., t.Any], ...],
 ],
 flat: t.Dict[
 t.Callable[..., t.Any],
 t.Tuple[
 "CallModel[..., t.Any]",
 t.Tuple[t.Callable[..., t.Any], ...],
 ],
 ],
) -> None:
 model, calls = items

 if model in collector:
 return

 if not calls:
 position = -1
 else:
 for i in calls:
 sub_model, _ = flat[i]
 if sub_model not in collector:
 _sort_dep(collector, flat[i], flat)

 position = max(collector.index(flat[i][0]) for i in calls)

 collector.insert(position + 1, model)


CUSTOM_ANNOTATIONS = (Depends, CustomField)


def build_call_model(
 call: t.Union[t.Callable[P, T], t.Callable[P, t.Awaitable[T]]],
 *,
 cast: bool = True,
 use_cache: bool = True,
 is_sync: t.Optional[bool] = None,
 extra_dependencies: t.Sequence[Depends] = (),
 pydantic_config: t.Optional[ConfigDict] = None,
) -> CallModel[P, T]:
 """Build a CallModel from a callable."""
 name = getattr(call, "__name__", type(call).__name__)

 is_call_async = is_coroutine_callable(call) or is_async_gen_callable(call)
 if is_sync is None:
 is_sync = not is_call_async
 else:
 assert not (is_sync and is_call_async), f"You cannot use async dependency `{name}` at sync main"

 typed_params, return_annotation = get_typed_signature(call)
 if (is_call_generator := is_gen_callable(call) or is_async_gen_callable(call)) and (
 return_args := get_args(return_annotation)
 ):
 return_annotation = return_args[0]

 class_fields: t.Dict[str, t.Tuple[t.Any, t.Any]] = {}
 dependencies: t.Dict[str, CallModel[..., t.Any]] = {}
 custom_fields: t.Dict[str, CustomField] = {}
 positional_args: t.List[str] = []
 keyword_args: t.List[str] = []
 var_positional_arg: t.Optional[str] = None
 var_keyword_arg: t.Optional[str] = None

 for param_name, param in typed_params.parameters.items():
 dep: t.Optional[Depends] = None
 custom: t.Optional[CustomField] = None

 if param.annotation is inspect.Parameter.empty:
 annotation = t.Any
 elif get_origin(param.annotation) is Annotated:
 annotated_args = get_args(param.annotation)
 type_annotation = annotated_args[0]

 custom_annotations = []
 regular_annotations = []
 for arg in annotated_args[1:]:
 if isinstance(arg, CUSTOM_ANNOTATIONS):
 custom_annotations.append(arg)
 else:
 regular_annotations.append(arg)

 assert len(custom_annotations) <= 1, (
 f"Cannot specify multiple `Annotated` Custom arguments for `{param_name}`!"
 )

 next_custom = next(iter(custom_annotations), None)
 if next_custom is not None:
 if isinstance(next_custom, Depends):
 dep = next_custom
 elif isinstance(next_custom, CustomField):
 custom = deepcopy(next_custom)
 else:
 raise AssertionError("unreachable")

 if regular_annotations:
 annotation = param.annotation
 else:
 annotation = type_annotation
 else:
 annotation = param.annotation
 else:
 annotation = param.annotation

 default: t.Any
 if param.kind == inspect.Parameter.VAR_POSITIONAL:
 default = ()
 var_positional_arg = param_name
 elif param.kind == inspect.Parameter.VAR_KEYWORD:
 default = {}
 var_keyword_arg = param_name
 elif param.default is inspect.Parameter.empty:
 default = Ellipsis
 else:
 default = param.default

 if isinstance(default, Depends):
 if dep:
 raise AssertionError(
 "You can not use both `Depends` with `Annotated` and a default",
 )
 dep, default = default, Ellipsis

 elif isinstance(default, CustomField):
 if custom:
 raise AssertionError(
 "You can not use both `CustomField` with `Annotated` and a default",
 )
 custom, default = default, Ellipsis

 else:
 class_fields[param_name] = (annotation, default)

 if dep:
 if not cast:
 dep.cast = False

 if isinstance(dep.dependency, solve_wrapper):
 dep.dependency = dep.dependency.call

 dependencies[param_name] = build_call_model(
 dep.dependency,
 cast=dep.cast,
 use_cache=dep.use_cache,
 is_sync=is_sync,
 pydantic_config=pydantic_config,
 )

 if dep.cast is True:
 class_fields[param_name] = (annotation, Ellipsis)

 keyword_args.append(param_name)

 elif custom:
 assert not (is_sync and is_coroutine_callable(custom.use)), (
 f"You cannot use async custom field `{type(custom).__name__}` at sync `{name}`"
 )

 custom.set_param_name(param_name)
 custom_fields[param_name] = custom

 if custom.cast is False:
 annotation = t.Any

 if custom.required:
 class_fields[param_name] = (annotation, default)
 else:
 class_fields[param_name] = class_fields.get(param_name, (t.Optional[annotation], None))

 keyword_args.append(param_name)

 else:
 if param.kind is param.KEYWORD_ONLY:
 keyword_args.append(param_name)
 elif param.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
 positional_args.append(param_name)

 func_model = create_model(
 name,
 __config__=get_config_base(pydantic_config),
 **class_fields,
 )

 response_model: t.Optional[t.Type[ResponseModel[T]]] = None
 if cast and return_annotation and return_annotation is not inspect.Parameter.empty:
 response_model = create_model(
 "ResponseModel",
 __config__=get_config_base(pydantic_config),
 response=(return_annotation, Ellipsis),
 )

 return CallModel(
 call=call,
 model=func_model,
 response_model=response_model,
 params=class_fields,
 cast=cast,
 use_cache=use_cache,
 is_async=is_call_async,
 is_generator=is_call_generator,
 dependencies=dependencies,
 custom_fields=custom_fields,
 positional_args=positional_args,
 keyword_args=keyword_args,
 var_positional_arg=var_positional_arg,
 var_keyword_arg=var_keyword_arg,
 extra_dependencies=[
 build_call_model(
 d.dependency,
 cast=d.cast,
 use_cache=d.use_cache,
 is_sync=is_sync,
 pydantic_config=pydantic_config,
 )
 for d in extra_dependencies
 ],
 )


class _InjectWrapper(t.Protocol[P, T]):
 def __call__(
 self,
 func: t.Callable[P, T],
 model: t.Optional[CallModel[P, T]] = None,
 ) -> t.Callable[P, T]: ...


@t.overload
def inject(
 func: None,
 *,
 cast: bool = True,
 extra_dependencies: t.Sequence[Depends] = (),
 pydantic_config: t.Optional[ConfigDict] = None,
 dependency_overrides_provider: t.Optional[t.Any] = dependency_provider,
 wrap_model: t.Callable[[CallModel[P, T]], CallModel[P, T]] = lambda x: x,
) -> _InjectWrapper[P, T]: ...


@t.overload
def inject(
 func: t.Callable[P, T],
 *,
 cast: bool = True,
 extra_dependencies: t.Sequence[Depends] = (),
 pydantic_config: t.Optional[ConfigDict] = None,
 dependency_overrides_provider: t.Optional[t.Any] = dependency_provider,
 wrap_model: t.Callable[[CallModel[P, T]], CallModel[P, T]] = lambda x: x,
) -> t.Callable[P, T]: ...


def inject(
 func: t.Optional[t.Callable[P, T]] = None,
 *,
 cast: bool = True,
 extra_dependencies: t.Sequence[Depends] = (),
 pydantic_config: t.Optional[ConfigDict] = None,
 dependency_overrides_provider: t.Optional[t.Any] = dependency_provider,
 wrap_model: t.Callable[[CallModel[P, T]], CallModel[P, T]] = lambda x: x,
) -> t.Union[t.Callable[P, T], _InjectWrapper[P, T]]:
 """Decorator to inject dependencies into a function."""
 decorator = _wrap_inject(
 dependency_overrides_provider=dependency_overrides_provider,
 wrap_model=wrap_model,
 extra_dependencies=extra_dependencies,
 cast=cast,
 pydantic_config=pydantic_config,
 )

 if func is None:
 return decorator
 else:
 return decorator(func)


def _wrap_inject(
 dependency_overrides_provider: t.Optional[t.Any],
 wrap_model: t.Callable[[CallModel[P, T]], CallModel[P, T]],
 extra_dependencies: t.Sequence[Depends],
 cast: bool,
 pydantic_config: t.Optional[ConfigDict],
) -> _InjectWrapper[P, T]:
 if (
 dependency_overrides_provider
 and getattr(dependency_overrides_provider, "dependency_overrides", None) is not None
 ):
 overrides = dependency_overrides_provider.dependency_overrides
 else:
 overrides = None

 def func_wrapper(
 func: t.Callable[P, T],
 model: t.Optional[CallModel[P, T]] = None,
 ) -> t.Callable[P, T]:
 if model is None:
 real_model = wrap_model(
 build_call_model(
 call=func,
 extra_dependencies=extra_dependencies,
 cast=cast,
 pydantic_config=pydantic_config,
 )
 )
 else:
 real_model = model

 if real_model.is_async:
 injected_wrapper: t.Callable[P, T]

 if real_model.is_generator:
 injected_wrapper = solve_wrapper(solve_async_gen, real_model, overrides)
 else:

 @wraps(func)
 async def injected_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
 async with AsyncExitStack() as stack:
 r = await real_model.asolve(
 *args,
 stack=stack,
 dependency_overrides=overrides,
 cache_dependencies={},
 nested=False,
 **kwargs,
 )
 return r

 raise AssertionError("unreachable")

 else:
 if real_model.is_generator:
 injected_wrapper = solve_wrapper(solve_gen, real_model, overrides)
 else:

 @wraps(func)
 def injected_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
 with ExitStack() as stack:
 r = real_model.solve(
 *args,
 stack=stack,
 dependency_overrides=overrides,
 cache_dependencies={},
 nested=False,
 **kwargs,
 )
 return r

 raise AssertionError("unreachable")

 return injected_wrapper

 return func_wrapper


class solve_async_gen:
 _iter: t.Optional[t.AsyncIterator[t.Any]] = None

 def __init__(
 self,
 model: "CallModel[..., t.Any]",
 overrides: t.Optional[t.Any],
 *args: t.Any,
 **kwargs: t.Any,
 ):
 self.call = model
 self.args = args
 self.kwargs = kwargs
 self.overrides = overrides

 def __aiter__(self) -> "solve_async_gen":
 self.stack = AsyncExitStack()
 return self

 async def __anext__(self) -> t.Any:
 if self._iter is None:
 stack = self.stack = AsyncExitStack()
 await self.stack.__aenter__()
 self._iter = t.cast(
 t.AsyncIterator[t.Any],
 (
 await self.call.asolve(
 *self.args,
 stack=stack,
 dependency_overrides=self.overrides,
 cache_dependencies={},
 nested=False,
 **self.kwargs,
 )
 ).__aiter__(),
 )

 try:
 r = await self._iter.__anext__()
 except StopAsyncIteration as e:
 await self.stack.__aexit__(None, None, None)
 raise e
 else:
 return r


class solve_gen:
 _iter: t.Optional[t.Iterator[t.Any]] = None

 def __init__(
 self,
 model: "CallModel[..., t.Any]",
 overrides: t.Optional[t.Any],
 *args: t.Any,
 **kwargs: t.Any,
 ):
 self.call = model
 self.args = args
 self.kwargs = kwargs
 self.overrides = overrides

 def __iter__(self) -> "solve_gen":
 self.stack = ExitStack()
 return self

 def __next__(self) -> t.Any:
 if self._iter is None:
 stack = self.stack = ExitStack()
 self.stack.__enter__()
 self._iter = t.cast(
 t.Iterator[t.Any],
 iter(
 self.call.solve(
 *self.args,
 stack=stack,
 dependency_overrides=self.overrides,
 cache_dependencies={},
 nested=False,
 **self.kwargs,
 )
 ),
 )

 try:
 r = next(self._iter)
 except StopIteration as e:
 self.stack.__exit__(None, None, None)
 raise e
 else:
 return r

Autostore: File Storage Made Simple

Sat, 14 Jun 2025 17:08:21 -0400

AutoStore provides a dictionary-like interface for reading and writing files with caching and different storage backends.

AutoStore eliminates the cognitive overhead of managing different file formats, letting you focus on your data and analysis rather than the mechanics of file I/O. It automatically handles file format detection, type inference, upload/download operations, and provides a clean, intuitive API for data persistence across local and cloud storage.

Table of contents:

Why Use AutoStore?
Getting Started
- Basic Usage
- Cloud Storage (S3)
Supported Data Types
Configuration Options
- S3StorageConfig
Advanced Features
Multiple Storage Backends
Performance Considerations
- Large File Handling
When to Use AutoStore
Comparison with Alternatives

Why Use AutoStore?

Simplicity: Store and retrieve data with dictionary syntax. No need to remember APIs for different file formats.
Caching: Caching system with configurable expiration reduces redundant downloads, especially for cloud storage.
Multiple Storage Backends: Seamlessly work with local files, S3, and other cloud storage services.
Type Detection: Automatically infers the best file format based on the data type.
Multiple Data Types: Built-in support for Polars DataFrames, JSON, CSV, images, PyTorch models, NumPy arrays, and more.
Extensible Architecture: Pluggable handler system for new data types and storage backends.
Performance Optimized: Upload/download operations with efficient handling of large files.
Type-Safe Configuration: Dataclass-based configuration with IDE support and validation.

Getting Started

AutoStore requires Python 3.10+ and can be installed via pip.

pip install autostore

Basic Usage

from autostore import AutoStore

store = AutoStore("./data")

# Write data - automatically saves with appropriate extensions
store["my_dataframe"] = df # ./data/my_dataframe.parquet
store["config"] = {"key": "value"} # ./data/config.json
store["logs"] = [{"event": "start"}] # ./data/logs.jsonl

# Read data
df = store["my_dataframe"] # Returns a Polars DataFrame
config = store["config"] # Returns a dict
logs = store["logs"] # Returns a list of dicts

Cloud Storage (S3)

from autostore import AutoStore
from autostore.s3 import S3Backend, S3StorageConfig

# Register S3 backend
AutoStore.register_backend("s3", S3Backend)

# Configure S3 with caching
s3_config = S3StorageConfig(
 region_name="us-east-1",
 cache_enabled=True,
 cache_expiry_hours=12,
 multipart_threshold=64 * 1024 * 1024 # 64MB
)

# Use S3 storage
store = AutoStore("s3://my-bucket/data/", config=s3_config)
store["experiment/results"] = {"accuracy": 0.95, "epochs": 100}
results = store["experiment/results"] # Uses cache on subsequent loads

Supported Data Types

Data Type	File Extension	Description	Library Required
Polars DataFrame/LazyFrame	`.parquet`, `.csv`	High-performance DataFrames	polars
Python dict/list	`.json`	Standard JSON serialization	built-in
List of dicts	`.jsonl`	JSON Lines format	built-in
Pydantic models	`.pydantic.json`	Structured data models	pydantic
Python dataclasses	`.dataclass.json`	Dataclass serialization	built-in
String data	`.txt`, `.html`, `.md`	Plain text files	built-in
NumPy arrays	`.npy`, `.npz`	Numerical data	numpy
SciPy sparse matrices	`.sparse`	Sparse matrix data	scipy
PyTorch tensors/models	`.pt`, `.pth`	Deep learning models	torch
PIL/Pillow images	`.png`, `.jpg`, etc.	Image data	Pillow
YAML data	`.yaml`, `.yml`	Human-readable config files	PyYAML
Any Python object	`.pkl`	Pickle fallback	built-in

Configuration Options

S3StorageConfig

from s3 import S3StorageConfig

config = S3StorageConfig(
 aws_access_key_id="your-key",
 aws_secret_access_key="your-secret",
 region_name="us-east-1",
 cache_enabled=True,
 cache_expiry_hours=12,
 multipart_threshold=64 * 1024 * 1024, # Files larger than this use multipart upload
 multipart_chunksize=16 * 1024 * 1024, # Chunk size for multipart uploads
 max_concurrency=10 # Maximum concurrent uploads/downloads
)

Advanced Features

Caching System

AutoStore includes an intelligent caching system that:

Stores frequently accessed files locally
Uses ETags for cache validation
Automatically expires old cache entries
Significantly improves performance for cloud storage

# Cache management
store.cleanup_cache() # Remove expired cache entries

# Check cache status
metadata = store.get_metadata("large_file")
print(f"File size: {metadata.size} bytes")
print(f"ETag: {metadata.etag}")

Custom Data Handlers

Add support for new data types by creating custom handlers:

from pathlib import Path
from autostore.autostore import DataHandler

class CustomLogHandler(DataHandler):
 def can_handle_extension(self, extension: str) -> bool:
 return extension.lower() == ".log"

 def can_handle_data(self, data) -> bool:
 return isinstance(data, list) and all(
 isinstance(item, dict) and "timestamp" in item
 for item in data
 )

 def read_from_file(self, file_path: Path, file_extension: str):
 logs = []
 with open(file_path, 'r') as f:
 for line in f:
 if line.strip():
 logs.append(json.loads(line))
 return logs

 def write_to_file(self, data, file_path: Path, file_extension: str):
 file_path.parent.mkdir(parents=True, exist_ok=True)
 with open(file_path, 'w') as f:
 for entry in data:
 f.write(json.dumps(entry) + '\n')

 @property
 def extensions(self):
 return [".log"]

 @property
 def priority(self):
 return 15

# Register the handler
store.register_handler(CustomLogHandler())

File Operations

# Check existence
if "config" in store:
 print("Config file exists")

# List all files
for key in store.keys():
 print(f"File: {key}")

# Get file metadata
metadata = store.get_metadata("large_dataset")
print(f"Size: {metadata.size} bytes")
print(f"Modified: {metadata.modified_time}")

# Copy and move files
store.copy("original", "backup")
store.move("temp_file", "permanent_file")

# Delete files
del store["old_data"]

Context Management

# Automatic cleanup of temporary files and cache
with AutoStore("./data", config=config) as store:
 store["data"] = large_dataset
 results = store["data"]
# Temporary files are automatically cleaned up here

Multiple Storage Backends

AutoStore supports pluggable storage backends:

# Local storage
local_store = AutoStore("./data")

# S3 storage
s3_store = AutoStore("s3://bucket/prefix/")

Performance Considerations

Large File Handling

AutoStore automatically optimizes for large files:

Multipart uploads/downloads for files > 64MB
Configurable chunk sizes and concurrency
Streaming operations to minimize memory usage

When to Use AutoStore

Choose AutoStore when you need:

Data science projects with mixed file types and cloud storage
Building data pipelines with heterogeneous data sources
Rapid prototyping where you don’t want to think about file formats
Consistent data access patterns across local and cloud environments
Performance optimization through intelligent caching
Easy extensibility for custom data types and storage backends
Type-safe configuration with dataclass-based settings

Don’t choose AutoStore when:

You need complex queries (use TinyDB or databases)
You only work with one data type consistently
You need zero dependencies (use Shelve)
You require advanced database features

Comparison with Alternatives

Feature	AutoStore	Shelve	DiskCache	TinyDB	PickleDB	SQLiteDict
Multi-format Support	✅ 12+ formats	❌ Pickle only	❌ Pickle only	❌ JSON only	❌ JSON only	❌ Pickle only
Auto Format Detection	✅ Smart inference	❌ Manual	❌ Manual	❌ Manual	❌ Manual	❌ Manual
Cloud Storage	✅ S3, extensible	❌ Local only	❌ Local only	❌ Local only	❌ Local only	❌ Local only
Intelligent Caching	✅ ETag-based	❌ None	✅ Advanced	❌ None	❌ None	❌ None
Type-Safe Config	✅ Dataclasses	❌ None	✅ Classes	❌ Dicts	❌ None	❌ None
Large File Handling	✅ Multipart	❌ Limited	✅ Good	❌ Limited	❌ Limited	❌ Limited
Extensibility	✅ Handler system	❌ Limited	❌ Limited	✅ Middleware	❌ Limited	❌ Limited
Performance	✅ Cached/Optimized	🔶 Medium	✅ Fast	🔶 Medium	🔶 Medium	🔶 Medium
Standard Library	❌ External	✅ Built-in	❌ External	❌ External	❌ External	❌ External

Presskit: Database-driven static site generator

Sat, 07 Jun 2025 00:00:00 +0000

Nearly all static site generators convert Markdown to HTML but aren’t good at generating multiple pages from database queries. Ideally we would store data in SQLite or Postgres databases and the static site generator would run queries against the database to create a page for each row in the result. Presskit was invented to do just that.

Presskit is a powerful static site generator that combines Markdown content with Jinja2 templating and database-driven page generation. Presskit lets you build dynamic static sites by connecting your content to SQLite databases and JSON data sources.

Table of contents:

Key Features
Installation
Quick Start
Basic Usage
Template Variables
Using Variables in Markdown
Data Sources and Queries
Generating Pages
Commands
- Build Commands
- Development
Advanced Configuration
- Full Configuration Example
- Custom Filters

Key Features

Jinja2 Templating: Use Jinja2 variables and logic in both Markdown content and HTML layouts
Database Integration: Load data from SQLite databases and JSON files
Dynamic Page Generation: Generate multiple pages automatically from SQLite query results
Structured Context: Access site metadata, build information, and data through a clean template context

Installation

pip install presskit

Or you can use Astral’s uv Python package manager to install Presskit as a self-contained tool so it can be run from the command line without needing to activate a virtual environment:

uv tool install presskit

Quick Start

Create a new site directory:

mkdir my-site
cd my-site

Create the basic structure:

my-site/
├── presskit.json # Configuration file
├── content/ # Markdown files
├── templates/ # HTML templates
└── public/ # Generated output (created automatically)

Build your site:

presskit build

Basic Usage

Writing Markdown Content

Create Markdown files in the content/ directory. Each file can include YAML front matter for metadata:

---
title: "Welcome to My Site"
description: "A brief introduction"
layout: page
---
# Welcome
This is my **awesome** site built with Presskit!

Creating HTML Templates

Templates go in the templates/ directory. Here’s a basic page.html template:

<!DOCTYPE html>
<html lang="{{ site.language }}">
<head>
 <meta charset="UTF-8">
 <title>{{ page.title or site.title }}</title>
 <meta name="description" content="{{ page.description or site.description }}">
</head>
<body>
 <header>
 <h1>{{ site.title }}</h1>
 </header>

 <main>
 {{ page.content }}
 </main>

 <footer>
 <p>&copy; {{ build.year }} {{ site.author }}</p>
 </footer>
</body>
</html>

Configuration

Create a presskit.json file to configure your site:

{
 "title": "My Awesome Site",
 "description": "A site built with Presskit",
 "author": "Your Name",
 "url": "https://mysite.com",
 "language": "en-US"
}

Template Variables

Presskit provides a structured context with the following variables available in all templates:

Site Variables (`site.*`)

site.title - Site title
site.description - Site description
site.author - Site author
site.url - Base site URL
site.version - Site version
site.language - Site language

Build Variables (`build.*`)

build.date - Build date (YYYY-MM-DD)
build.year - Build year
build.timestamp - Full build timestamp
build.iso_date - Build date in ISO format

Page Variables (`page.*`)

page.filename - Page filename without extension
page.filepath - Full file path
page.path - Clean URL path
page.layout - Template layout name
page.content - Processed HTML content (in templates)
page.title - Page title from front matter
page.description - Page description from front matter

Data Variables (`data.*`)

data.queries - Results from named queries
data.sources - JSON data sources
data.page_queries - Page-specific query results

Plus any custom variables from your front matter are available at the top level.

Using Variables in Markdown

You can use Jinja2 templating directly in your Markdown content:

---
title: About
category: personal
---
# About {{ site.author }}
This site was built on {{ build.date }} and is currently version {{ site.version }}.
{% if category == "personal" %}
This is a personal page about {{ site.author }}.
{% endif %}

Data Sources and Queries

Presskit’s data integration feature allows you to connect your static site to data sources, enabling content generation while maintaining the performance benefits of static sites. This powerful feature bridges the gap between static and dynamic websites.

This enables data-driven pages that display statistics, reports, or any structured data. Ideal for portfolios showcasing project metrics, business dashboards, or documentation sites pulling from APIs.

This encourages separation of concerns where you keep your content in databases where it can be easily edited, queried, and managed, while your site structure remains in version control.

Configuring Data Sources

Add data sources to your presskit.json:

{
 "title": "My Blog",
 "sources": {
 "blog_db": {
 "type": "sqlite",
 "path": "data/blog.db"
 },
 "config": {
 "type": "json",
 "path": "data/site-config.json"
 }
 },
 "default_source": "blog_db"
}

Adding Queries

Define queries to load data from your sources:

{
 "sources": {
 "blog_db": {
 "type": "sqlite",
 "path": "data/blog.db"
 }
 },
 "queries": [
 {
 "name": "recent_posts",
 "source": "blog_db",
 "query": "SELECT title, slug, date, excerpt FROM posts ORDER BY date DESC LIMIT 5"
 },
 {
 "name": "categories",
 "source": "blog_db",
 "query": "SELECT name, slug, COUNT(*) as post_count FROM categories JOIN posts ON categories.id = posts.category_id GROUP BY categories.id"
 }
 ]
}

Using Query Data in Templates

Access query results through the data.queries object:

<section class="recent-posts">
 <h2>Recent Posts</h2>
 {% for post in data.queries.recent_posts %}
 <article>
 <h3><a href="/posts/{{ post.slug }}">{{ post.title }}</a></h3>
 <time>{{ post.date | date_format('%B %d, %Y') }}</time>
 <p>{{ post.excerpt }}</p>
 </article>
 {% endfor %}
</section>

<aside class="categories">
 <h3>Categories</h3>
 <ul>
 {% for category in data.queries.categories %}
 <li><a href="/category/{{ category.slug }}">{{ category.name }} ({{ category.post_count }})</a></li>
 {% endfor %}
 </ul>
</aside>

Page-Level Queries

You can also define queries in individual Markdown files:

---
title: "Author Profile"
queries:
 author_posts:
 source: "blog_db"
 query: "SELECT title, slug, date FROM posts WHERE author_id = {{ author_id }} ORDER BY date DESC"
variables:
 author_id: 123
---

# {{ author.name }}

## Recent Posts by This Author

{% for post in data.page_queries.author_posts %}
- [{{ post.title }}](/posts/{{ post.slug }}) - {{ post.date | date_format('%Y-%m-%d') }}
{% endfor %}

The above example shows how to define a query that fetches posts by a specific author using the author_id variable.

Generating Pages

The most powerful feature of Presskit is generating multiple pages from database queries.

Generator Queries

Mark a query as a generator to create multiple pages:

{
 "queries": [
 {
 "name": "blog_posts",
 "source": "blog_db",
 "query": "SELECT title, slug, content, date, author FROM posts WHERE published = 1",
 "generator": true,
 "template": "post",
 "output_path": "posts/#{slug}"
 }
 ]
}

Generator Configuration

generator: true - Marks this as a page generator
template - Template to use for generated pages
output_path - Path pattern with placeholders like #{field_name}

Creating Generator Templates

Create a template for your generated pages (templates/post.html):

<!DOCTYPE html>
<html>
<head>
 <title>{{ title }} | {{ site.title }}</title>
</head>
<body>
 <article>
 <h1>{{ title }}</h1>
 <time>{{ date | date_format('%B %d, %Y') }}</time>
 <div class="content">
 {{ content | safe }}
 </div>
 <p>By {{ author }}</p>
 </article>

 <nav>
 <a href="/">← Back to Home</a>
 </nav>
</body>
</html>

Nested Queries

You can create parent-child query relationships:

{
 "queries": [
 {
 "name": "authors",
 "source": "blog_db",
 "query": "SELECT id, name, bio, slug FROM authors"
 },
 {
 "name": "authors.posts",
 "source": "blog_db",
 "query": "SELECT title, slug, date FROM posts WHERE author_id = {{ id }} ORDER BY date DESC"
 }
 ]
}

Access nested data in templates:

{% for author in data.queries.authors %}
<div class="author">
 <h2>{{ author.name }}</h2>
 <p>{{ author.bio }}</p>

 <h3>Posts by {{ author.name }}</h3>
 {% for post in author.posts %}
 <p><a href="/posts/{{ post.slug }}">{{ post.title }}</a> - {{ post.date }}</p>
 {% endfor %}
</div>
{% endfor %}

Commands

Build Commands

# Build entire site
presskit build

# Build specific file
presskit build content/about.md

# Execute queries and cache results
presskit data

# Generate pages from generator queries 
presskit generate

# Check query cache status
presskit status

Development

# Start development server
presskit server

# Clean build artifacts
presskit clean

Advanced Configuration

Full Configuration Example

{
 "title": "My Blog",
 "description": "A blog about web development",
 "author": "Jane Developer",
 "url": "https://myblog.dev",
 "version": "2.1.0",
 "language": "en-US",

 "content_dir": "content",
 "templates_dir": "templates",
 "output_dir": "public",
 "cache_dir": ".cache",

 "default_template": "page",
 "markdown_extension": "md",
 "workers": 8,

 "server_host": "0.0.0.0",
 "server_port": 8000,

 "sources": {
 "blog_db": {
 "type": "sqlite",
 "path": "data/blog.sqlite3"
 },
 "config": {
 "type": "json",
 "path": "data/config.json"
 }
 },

 "default_source": "blog_db",

 "variables": {
 "environment": "production",
 "analytics_id": "GA-XXXXX"
 },

 "queries": [
 {
 "name": "posts",
 "source": "blog_db",
 "query": "SELECT * FROM posts WHERE status = 'published' ORDER BY date DESC",
 "generator": true,
 "template": "post",
 "output_path": "blog/#{slug}"
 },
 {
 "name": "recent_posts",
 "source": "blog_db",
 "query": "SELECT title, slug, excerpt, date FROM posts WHERE status = 'published' ORDER BY date DESC LIMIT 5"
 }
 ]
}

Custom Filters

Presskit includes useful Jinja2 filters:

date_format(format)
Format dates (e.g., {{ date | date_format('%B %d, %Y') }})

Google Chrome On-Device Embedding Model

Tue, 13 May 2025 00:00:00 +0000

Google Chrome bundles a text embedding model used to cluster browsing history as part of the Topics API and for semantic search. They also ship a number of other models with Chrome.

First I had to track down these on-device models. I started at the usual place where apps store their application data in the ~/Library/Application Support/ folder on macOS. Searching for find ~/Library/Application\ Support/ -maxdepth 4 'optimization*' returns the the folder I was looking for: ~/Library/Application Support/Google/Chrome/optimization_guide_model_store. The Chromium source code says this about optimization guide: “The optimization guide component contains code for processing hints and machine learning models received from the remote Chrome Optimization Guide Service”.

Models are stored as tflite files, which is rebranded as LiteRT:

LiteRT (short for Lite Runtime), formerly known as TensorFlow Lite, is Google’s high-performance runtime for on-device AI. LiteRT is also the format Google uses to ship their Gemma3N edge-device models.

$ cd ~/Library/Application\ Support/Google/Chrome/optimization_guide_model_store
$ tree -L 4 .
.
├── 13
│   └── E6DC4029A1E4B4C1
│   └── 205CA176C885321E
│   ├── model-info.pb
│   └── model.tflite
├── 15
│   └── E6DC4029A1E4B4C1
│   └── 255B83C178FA9DD9
│   ├── model-info.pb
│   ├── model.tflite
│   ├── override_list.pb.gz
│   └── VERSION.txt
├── 2
│   └── E6DC4029A1E4B4C1
│   └── 2DFDB6405E512759
│   ├── model-info.pb
│   └── model.tflite
├── 20
│   └── E6DC4029A1E4B4C1
│   └── 91EF641BEE15B40C
│   ├── model-info.pb
│   └── model.tflite
├── 24
│   └── E6DC4029A1E4B4C1
│   └── 1B8C0D25285420AB
│   ├── enus_denylist_encoded_241007.txt
│   ├── model-info.pb
│   ├── model.tflite
│   └── vocab_en-us.txt
├── 25
│   └── E6DC4029A1E4B4C1
│   └── C278361C6A5A6107
│   ├── model-info.pb
│   ├── model.tflite
│   └── visual_model_desktop.tflite
├── 26
│   └── E6DC4029A1E4B4C1
│   └── 141FCE0CF6807549
│   ├── model-info.pb
│   └── model.tflite
├── 43
│   └── E6DC4029A1E4B4C1
│   └── E234446CB5BACE99
│   ├── model-info.pb
│   ├── model.tflite
│   └── sentencepiece.model
├── 45
│   └── E6DC4029A1E4B4C1
│   └── 063B3FABDDE10CE8
│   ├── model-info.pb
│   └── model.tflite
└── 9
 └── E6DC4029A1E4B4C1
 └── B5ECF67C32B2BD47
 ├── model-info.pb
 └── model.tflite

31 directories, 26 files

One of the models (255B83C178FA9DD9) is the “Browsing Topics Privacy Sandbox feature” which maps recent browsing history to a set of interest-based categories to serve relevant ads. The Topics API is a replacement for the FLOC proposal. The version.txt refers to a taxonomy, which presumably is related to the interest-based categories.

The visual_model_desktop.tflite file appears to be part of a phishing classifier based on the Chromium source code.

There is a sentencepiece.model file which is a text-tokenizer. The accompanying model.tflite file is the largest at 107 MB.

$ find . -type f -name "*model.tflite" -print0 | xargs -0 du -h

64K ./20/E6DC4029A1E4B4C1/91EF641BEE15B40C/model.tflite
16K ./9/E6DC4029A1E4B4C1/B5ECF67C32B2BD47/model.tflite
132K ./45/E6DC4029A1E4B4C1/063B3FABDDE10CE8/model.tflite
8.0K ./26/E6DC4029A1E4B4C1/141FCE0CF6807549/model.tflite
107M ./43/E6DC4029A1E4B4C1/E234446CB5BACE99/model.tflite <---
4.4M ./24/E6DC4029A1E4B4C1/1B8C0D25285420AB/model.tflite
2.6M ./15/E6DC4029A1E4B4C1/255B83C178FA9DD9/model.tflite
384K ./2/E6DC4029A1E4B4C1/2DFDB6405E512759/model.tflite
1.2M ./13/E6DC4029A1E4B4C1/205CA176C885321E/model.tflite
180K ./25/E6DC4029A1E4B4C1/C278361C6A5A6107/model.tflite

Let’s load up the sentencepiece tokenizer and the model and get embeddings.

import numpy as np
import tensorflow as tf
import sentencepiece as spm

tokenizer = spm.SentencePieceProcessor(model_file='sentencepiece.model')
interpreter = tf.lite.Interpreter(model_path='model.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

def get_embedding(
 interpreter: tf.lite.Interpreter,
 tokenizer: spm.SentencePieceProcessor,
 text: str,
) -> np.ndarray:
 """
 Embedding vector for a given text. Max token length is 64.

 Args:
 interpreter: TFLite interpreter.
 tokenizer: SentencePiece tokenizer.
 text: Text to be tokenized and embedded.

 Returns:
 np.ndarray: Embedding vector of shape (768,).

 Example:
 >>> tokenizer = spm.SentencePieceProcessor(model_file="sentencepiece.model")
 >>> interpreter = tf.lite.Interpreter(model_path="model.tflite")
 >>> get_embedding(interpreter, tokenizer, "New York") # (768,)
 """
 input_shape = input_details[0]["shape"]
 seq_len = input_shape[1] if len(input_shape) > 1 else 64
 tokens = tokenizer.encode(text, out_type=int)
 tokens = np.pad(
 tokens[:seq_len],
 (0, max(0, seq_len - len(tokens))),
 "constant"
 )[np.newaxis, :] # (1, 64)
 interpreter.set_tensor(input_details[0]["index"], tokens.astype(np.int32))
 interpreter.invoke()
 embedding = interpreter.get_tensor(output_details[0]["index"]) # (1, 768)
 return embedding[0] # (768,)

text = "New York"
embedding = get_embedding(interpreter, tokenizer, text)
# [-2.49071661e-02 2.41535041e-03 -1.51733104e-02 -1.12882648e-02...]

This model has a max sequence length of 64 inpt tokens and outputs a 768 dimension vector. That is a small input size of around 200-220 characters.

Automatic statistics selection

Sat, 10 May 2025 00:00:00 +0000

Given a dataset with mixed variable types, how do you decide which statistical test to run? We can reduce this down to three things: the type of the variable being tested (numeric or categorical), the type of the covariate (numeric, categorical, or none), and how many categories the covariate has.

The following decision table maps these inputs to the appropriate test. Each function takes a Polars DataFrame, runs the test, and returns a structured result with a test statistic, critical value, p-value, and a plain-language conclusion.

Variable	Covariate	NumCategories	Type	Test	StatementExample
Numeric	Categorical	>2	Medians	Kruskal-Wallis	“Median does not vary across values of .”
Numeric	Categorical	2	Medians	Mann-Whitney Test	“Median varies across values of .”
Numeric	Categorical	>2	Means	ANOVA	“Mean does not vary across values of .”
Numeric	Categorical	2	Means	T-test	“Mean varies across values of .”
Categorical	Numeric	2+	Chi-squared	Chi-squared test	“ and quantile of are not independent”
Numeric	None	1	Distribution	Kolmogorov-Smirnov	“ does not depart from a uniform distribution.”
Numeric	Numeric	6	Distribution	Kolmogorov-Smirnov	“The distribution of varies across quantile of .”
Numeric	Categorical	2+	Distribution	Kolmogorov-Smirnov	“The distribution of does not vary across values of .”
Numeric	Numeric	2	Correlation	Pearson Correlation	“A positive correlation exists between and .”
Numeric	Categorical	2+	Rank Correlation	Spearman Correlation	“A negative correlation exists between and .”

In the case where variable is a categorical and covariate is numeric, the test is a Chi-squared test of independence. The numeric variable is binned into quantiles, and the Chi-squared test is performed on the contingency table of the binned variable and the categorical covariate. When the categorial has only 2 levels and the covariate is only 2 levels, then we have a special case of the Chi-squared test of independence, which is the Fisher’s exact test and we report the z-score and p-value of the Fisher’s exact test. Otherwise, we report the Chi-squared test statistic and p-value.

The Kolmogorov-Smirnov test can be used in three ways:

Single variable: Tests whether a numeric variable follows a uniform distribution (normalized to [0,1] range).
With numeric covariate: Uses two-sample KS tests to compare distributions across 6 quantiles of the covariate.
With categorical covariate: Uses two-sample KS tests to compare distributions across category groups.

The Pearson Correlation test examines the linear relationship between two numeric variables. It returns the correlation coefficient as the test statistic along with degrees of freedom (n-2). The test determines if a significant positive or negative correlation exists between the variables.

The Spearman Correlation test examines the monotonic relationship between a numeric variable and categorical covariate. The categorical variable is converted to numeric ranks (sorted alphabetically) and then Spearman’s rank correlation is computed. This test is useful for detecting ordinal relationships between numeric and categorical variables.

We can turn these rules into a function that automatically selects and runs the appropriate test based on the variable types and covariate characteristics. The function will return a structured result with the test statistic, critical value, p-value, and a plain-language conclusion.

import numpy as np # type: ignore
import polars as pl
from scipy import stats # type: ignore
from pydantic import BaseModel
from typing import List, Union, Optional


class MeanDiffBootstrapResult(BaseModel):
 mean: float
 ci_lb: float
 ci_ub: float


class EffectSizeResult(BaseModel):
 standardized_difference: float
 ci_lower: float
 ci_upper: float


class EffectSizeCategoricalResult(BaseModel):
 standardized_difference: Optional[float]
 ci_lower: Optional[float]
 ci_upper: Optional[float]


class StandardizedDifferenceVariable(BaseModel):
 variable: str
 type: str
 standardized_difference: Optional[float]
 ci_lower: Optional[float]
 ci_upper: Optional[float]


class StandardizedDifferenceResult(BaseModel):
 variables: List[StandardizedDifferenceVariable]


class StatisticalTestResult(BaseModel):
 test_name: str
 test_statistic: float
 critical_value: float
 significant: bool
 p_value: float
 conclusion: str
 degrees_of_freedom: Optional[int] = None
 contingency_table: Optional[dict] = None


def bin_numeric_column(df: pl.DataFrame, column: str, n_bins: int = 6, method: str = "quantile") -> pl.DataFrame:
 """
 Bin a numeric column in a dataframe using either quantile or equidistant binning.

 For quantile binning, extreme outliers (below 1st percentile and above 99th percentile)
 are removed when calculating quantile thresholds, but all data points are included in
 the final bins. Outliers are placed in the appropriate edge bins.

 Args:
 df (pl.DataFrame): Input dataframe
 column (str): Name of the numeric column to bin
 n_bins (int): Number of bins to create (default: 4)
 method (str): Binning method - "quantile" or "equidistant" (default: "quantile")

 Returns:
 pl.DataFrame: Dataframe with an additional column "{column}_bins" containing bin labels
 """
 if method == "quantile":
 # Use quantile-based binning with outlier winsorization
 # Get non-null values for outlier detection
 non_null_values = df.filter(pl.col(column).is_not_null())[column].to_numpy()

 if len(non_null_values) == 0:
 # Handle all-null column
 return df.with_columns(pl.lit(None, dtype=pl.Utf8).alias(f"{column}_bins"))

 if len(set(non_null_values)) == 1:
 # Handle constant column case
 return df.with_columns(pl.lit("Q1").alias(f"{column}_bins"))

 # Remove extreme outliers for quantile calculation (winsorization)
 p1 = np.percentile(non_null_values, 1)
 p99 = np.percentile(non_null_values, 99)

 # Filter out extreme outliers for quantile calculation only
 winsorized_values = non_null_values[(non_null_values >= p1) & (non_null_values <= p99)]

 if len(winsorized_values) < n_bins:
 # If too few values after winsorization, use all values
 winsorized_values = non_null_values

 # Calculate quantile probs based on winsorized data, then apply to all data
 quantile_probs = [i / n_bins for i in range(n_bins + 1)]
 quantile_values = [np.percentile(winsorized_values, q * 100) for q in quantile_probs]

 # Adjust edges to include all original data (including outliers)
 quantile_values[0] = non_null_values.min() - 1e-10
 quantile_values[-1] = non_null_values.max() + 1e-10

 # Use cut with manual binning
 def assign_bin(value):
 if value is None:
 return None
 for i in range(len(quantile_values) - 1):
 if quantile_values[i] <= value < quantile_values[i + 1]:
 return f"Q{i + 1}"
 # Handle edge case for maximum value
 return f"Q{n_bins}"

 # Apply binning using map_elements
 return df.with_columns(pl.col(column).map_elements(assign_bin, return_dtype=pl.Utf8).alias(f"{column}_bins"))

 elif method == "equidistant":
 # For n_bins output bins, create n_bins - 1 internal break points (excluding the min/max endpoints)
 # This allows Polars to automatically create exactly n_bins intervals: (-inf, break1], (break1, break2], ..., (breakN-1, inf]
 # Provide n_bins labels to match these intervals

 # Use equidistant binning
 col_min = df[column].min()
 col_max = df[column].max()

 # Handle null cases
 if col_min is None or col_max is None:
 return df.with_columns(pl.lit(None, dtype=pl.Utf8).alias(f"{column}_bins"))

 if col_min == col_max:
 # Handle constant column case
 return df.with_columns(pl.lit("Bin1").alias(f"{column}_bins"))

 # Create equidistant bins
 # For n_bins, we need n_bins-1 internal break points (excluding min/max)
 # Polars creates bins: (-inf, break1], (break1, break2], ..., (breakN-1, inf]
 if n_bins == 1:
 # Special case: only one bin
 return df.with_columns(pl.lit("Bin1").alias(f"{column}_bins"))

 bin_edges = np.linspace(float(col_min), float(col_max), n_bins + 1)[1:-1] # exclude endpoints
 bin_labels = [f"Bin{i + 1}" for i in range(n_bins)]

 # Use cut for equidistant binning
 return df.with_columns(pl.col(column).cut(bin_edges.tolist(), labels=bin_labels).alias(f"{column}_bins"))
 else:
 raise ValueError(f"Unknown binning method: {method}. Use 'quantile' or 'equidistant'")


def mean_diff(group_a, group_b) -> float:
 return np.mean(group_a) - np.mean(group_b)


def drop_outliers_array(arr, quantiles):
 lower = np.percentile(arr, quantiles[0] * 100)
 upper = np.percentile(arr, quantiles[1] * 100)
 # Drop outliers outside the quantiles
 mask = (arr >= lower) & (arr <= upper)
 return arr[mask]


def bootstrap_mean_diff(
 df: Union[pl.DataFrame, pl.LazyFrame],
 metric_col: str,
 intervention_col: str,
 n_resamples: int = 1000,
 drop_outliers: bool = False,
 outlier_quantiles: List[float] = [0.01, 0.99],
) -> MeanDiffBootstrapResult:
 """Calculate the mean difference between two groups using bootstrapping.

 Args:
 df (Union[pl.DataFrame, pl.LazyFrame]): The input DataFrame containing the data.
 metric_col (str): The name of the column containing the metric to analyze.
 intervention_col (str): The name of the column indicating the intervention group (1 for treatment, 0 for control).
 n_resamples (int): The number of bootstrap resamples to perform. Default is 1000.
 drop_outliers (bool): Whether to drop outliers outside the specified quantiles. Default is False.
 outlier_quantiles (List[float]): The quantiles to use for outlier removal. Default is [0.01, 0.99].
 Values outside these quantiles will be dropped from the analysis.
 """
 # Separate the data into two groups based on the 'intervention' column
 # Select only needed columns and collect for analysis
 if isinstance(df, pl.LazyFrame):
 dx = df.select([metric_col, intervention_col]).drop_nulls().collect()
 else:
 dx = df.lazy().select([metric_col, intervention_col]).drop_nulls().collect()

 # Group_a is the treatment group (intervention_col == 1)
 # Group_b is the control group (intervention_col == 0)
 if drop_outliers:
 # Apply outlier removal to each group after splitting
 group_a = dx.filter(pl.col(intervention_col) == 1)[metric_col].to_numpy()
 group_b = dx.filter(pl.col(intervention_col) == 0)[metric_col].to_numpy()
 group_a = drop_outliers_array(group_a, outlier_quantiles)
 group_b = drop_outliers_array(group_b, outlier_quantiles)
 else:
 group_a = dx.filter(pl.col(intervention_col) == 1)[metric_col].to_numpy()
 group_b = dx.filter(pl.col(intervention_col) == 0)[metric_col].to_numpy()
 # Bootstrap resampling
 res = stats.bootstrap(
 (group_a, group_b),
 statistic=mean_diff,
 n_resamples=n_resamples,
 method="percentile",
 )
 ci = res.confidence_interval
 mu = float(np.mean(res.bootstrap_distribution))
 return MeanDiffBootstrapResult(mean=mu, ci_lb=ci.low, ci_ub=ci.high)


def effect_size_continuous(
 df: Union[pl.DataFrame, pl.LazyFrame], group_col: str, var_col: str, coverage: float = 0.95, decimals: int = 3
) -> EffectSizeResult:
 """Effect size for continuous variables.

 Args:
 df (pl.DataFrame): Dataframe
 group_col (str): Group column name with treatment and control assignments
 var_col (str): Variable column name
 coverage (float): Coverage of the confidence interval
 decimals (int): Number of decimals to round the effect size

 Returns:
 EffectSizeResult: Effect size and confidence interval

 Examples:

 ```python
 df = pl.DataFrame(
 {
 "group": ["A", "B", "A", "B", "A", "B", "A", "B"],
 "var": [0.1, 0.2, 0.13, 0.4, 0.25, 0.6, 0.17, 0.8],
 }
 )
 effect_size_continuous(df, "group", "var") # EffectSizeResult(standardized_difference=1.793, ci_lower=0.152, ci_upper=3.434)
 ```

 """
 import math
 import scipy.stats as stats

 # Select only needed columns and collect for analysis
 if isinstance(df, pl.LazyFrame):
 df_subset = df.select([group_col, var_col]).drop_nulls().collect()
 else:
 df_subset = df.select([group_col, var_col]).drop_nulls()

 assert df_subset[group_col].n_unique() == 2, "Only 2 groups allowed"

 # Calculate means and variances by group
 group_stats = (
 df_subset.group_by(group_col)
 .agg(
 [
 pl.col(var_col).mean().alias("mean"),
 pl.col(var_col).var(ddof=1).alias("var"),
 pl.col(var_col).count().alias("count"),
 ]
 )
 .sort(group_col)
 )

 m = group_stats["mean"].to_numpy()
 v = group_stats["var"].to_numpy()
 n = group_stats["count"].to_numpy()

 stdiff = (m[1] - m[0]) / math.sqrt(v.sum() / 2)
 stdiff = round(stdiff, decimals)

 # compute confidence interval
 # Number of observations in group 1, group 0, and total
 n0, n1 = n[0], n[1]
 total = n0 + n1
 # Computing the corresponding value from the standard Normal for specified CI coverage
 percentile = 1 - ((1 - coverage) / 2)
 zscore = stats.norm.ppf(percentile)
 # Computing the standard deviation
 deviation = np.sqrt((total / (n0 * n1)) + ((stdiff**2) / (2 * total)))
 # Constructing the CIs using the Z-score and standard deviation
 lower_ci = stdiff - zscore * deviation
 upper_ci = stdiff + zscore * deviation
 lower_ci = round(lower_ci, decimals)
 upper_ci = round(upper_ci, decimals)
 return EffectSizeResult(standardized_difference=stdiff, ci_lower=lower_ci, ci_upper=upper_ci)


def effect_size_categorical(
 df: Union[pl.DataFrame, pl.LazyFrame], group_col: str, var_col: str, coverage: float = 0.95, decimals: int = 3
) -> EffectSizeCategoricalResult:
 """Effect size for binary categorical variables.

 Args:
 df (pl.DataFrame): Dataframe
 group_col (str): Group column name with treatment and control assignments
 var_col (str): Variable column name, binary (2 levels) or categorical (>2 levels)
 coverage (float): Coverage of the confidence interval
 decimals (int): Number of decimals to round the effect size

 Returns:
 EffectSizeCategoricalResult: Effect size for each level

 Examples:

 ```python
 df = pl.DataFrame({
 "group": ["T", "C", "T", "C", "T", "C"],
 "var": ["A", "B", "A", "A", "B", "C"],
 })
 effect_size_categorical(df, "group", "var") # EffectSizeCategoricalResult(standardized_difference=1.069, ci_lower=-0.642, ci_upper=2.78)
 ```
 """
 import scipy.stats as stats

 # Select only needed columns and collect for analysis
 if isinstance(df, pl.LazyFrame):
 df = df.select([group_col, var_col]).drop_nulls().collect()
 else:
 df = df.select([group_col, var_col]).drop_nulls()

 # Calculate value counts for each group-variable combination
 value_counts = df.group_by([group_col, var_col]).agg(pl.len().alias("count"))

 # Calculate total counts per group to compute proportions
 group_totals = df.group_by(group_col).agg(pl.len().alias("total"))

 # Join to get proportions
 probs_df = value_counts.join(group_totals, on=group_col).with_columns(
 (pl.col("count") / pl.col("total")).alias("prob")
 )

 # Pivot to get groups as rows and variable levels as columns
 # Get unique values for both dimensions
 groups = df[group_col].unique().sort().to_list()
 levels = df[var_col].unique().sort().to_list()

 # Create a matrix of probabilities
 prob_matrix = np.zeros((len(groups), len(levels)))

 for i, group in enumerate(groups):
 for j, level in enumerate(levels):
 prob_row = probs_df.filter((pl.col(group_col) == group) & (pl.col(var_col) == level))
 if len(prob_row) > 0:
 prob_matrix[i, j] = prob_row["prob"][0]

 group0 = prob_matrix[0, :]
 group1 = prob_matrix[1, :]
 # prob_matrix.shape = (levels, 2)
 prob_matrix_t = prob_matrix.transpose()
 # Computing the probability difference between group 1 and group 0
 # Dropping the 1st difference as there are n-1 degrees of freedom
 prob_difference = np.subtract(group1, group0)
 prob_difference = np.delete(prob_difference, (0), axis=0)

 # Check if we have any probability differences left after dropping the first
 if len(prob_difference) == 0:
 # Handle case with only 2 levels (binary categorical)
 # Use simple difference between the probabilities for the second level
 prob_diff_scalar = group1[1] - group0[1]
 # For binary case, use simple variance formula
 var_group0 = group0[1] * (1 - group0[1])
 var_group1 = group1[1] * (1 - group1[1])
 pooled_var = (var_group0 + var_group1) / 2

 if pooled_var == 0:
 # Perfect separation case
 return EffectSizeCategoricalResult(standardized_difference=None, ci_lower=None, ci_upper=None)

 stdiff = abs(prob_diff_scalar) / np.sqrt(pooled_var)
 stdiff = round(stdiff, decimals)

 # compute confidence interval for binary case
 group_sizes = df.group_by(group_col).agg(pl.len().alias("size")).sort(group_col)
 n0, n1 = group_sizes["size"].to_list()
 total = n0 + n1
 percentile = 1 - ((1 - coverage) / 2)
 zscore = stats.norm.ppf(percentile)
 deviation = np.sqrt((total / (n0 * n1)) + ((stdiff**2) / (2 * total)))
 lower_ci = stdiff - zscore * deviation
 upper_ci = stdiff + zscore * deviation
 lower_ci = round(lower_ci, decimals)
 upper_ci = round(upper_ci, decimals)

 return EffectSizeCategoricalResult(standardized_difference=stdiff, ci_lower=lower_ci, ci_upper=upper_ci)

 # Computing the covariance matrix
 levels_count = prob_matrix_t.shape[0]
 covariance = np.zeros(shape=(levels_count, levels_count))
 for row in range(levels_count):
 for col in range(levels_count):
 if row == col:
 covariance[row][col] = (
 prob_matrix_t[row][0] * (1 - prob_matrix_t[row][0])
 + prob_matrix_t[row][1] * (1 - prob_matrix_t[row][1])
 ) / 2
 else:
 covariance[row][col] = (
 -(prob_matrix_t[row][0] * prob_matrix_t[col][0] + prob_matrix_t[row][1] * prob_matrix_t[col][1]) / 2
 )
 # Dropping the 1st line and row as there are n-1 degrees of freedom
 # Computing the inverse of the covariance matrix
 covariance = np.delete(covariance, (0), axis=0)
 covariance = np.delete(covariance, (0), axis=1)
 try:
 inverse = np.linalg.inv(covariance)
 except np.linalg.LinAlgError:
 return EffectSizeCategoricalResult(standardized_difference=None, ci_lower=None, ci_upper=None)
 # Computing the standardized difference (using Mahalanobis distance)
 stdiff = np.sqrt(np.linalg.multi_dot([prob_difference.T, inverse, prob_difference]))
 stdiff = round(stdiff, decimals)
 # compute confidence interval
 # Number of observations in group 1, group 0, and total
 group_sizes = df.group_by(group_col).agg(pl.len().alias("size")).sort(group_col)
 n0, n1 = group_sizes["size"].to_list()
 total = n0 + n1
 # Computing the corresponding value from the standard Normal for specified CI coverage
 percentile = 1 - ((1 - coverage) / 2)
 zscore = stats.norm.ppf(percentile)
 # Computing the standard deviation
 deviation = np.sqrt((total / (n0 * n1)) + ((stdiff**2) / (2 * total)))
 # Constructing the CIs using the Z-score and standard deviation
 lower_ci = stdiff - zscore * deviation
 upper_ci = stdiff + zscore * deviation
 lower_ci = round(lower_ci, decimals)
 upper_ci = round(upper_ci, decimals)
 return EffectSizeCategoricalResult(standardized_difference=stdiff, ci_lower=lower_ci, ci_upper=upper_ci)


def standardized_difference(
 df: Union[pl.DataFrame, pl.LazyFrame],
 intervention: str,
 categorical: Optional[Union[str, List[str]]] = None,
 continuous: Optional[Union[str, List[str]]] = None,
 coverage: float = 0.95,
 decimals: int = 3,
) -> pl.DataFrame:
 """Effect size for binary categorical variables and continuous variables.

 Args:
 df (pl.DataFrame): Dataframe
 intervention (str): Intervention column name with treatment and control assignments
 categorical (str or list): Categorical variable column name(s)
 continuous (str or list): Continuous variable column name(s)
 coverage (float): Coverage of the confidence interval, default 0.95
 decimals (int): Number of decimals to round the effect size, default 3

 Returns:
 (pl.DataFrame): Dataframe with effect size for each level

 Examples:
 ```python
 df = pl.DataFrame({
 "intervention": ["A", "B", "A", "B", "A", "B"],
 "categorical_var": ["X", "Y", "Y", "Y", "X", "X"], # Now both groups have both X and Y
 "continuous_var": [1.2, 2.3, 1.5, 2.8, 1.7, 2.9],
 })
 standardized_difference(df, "intervention", categorical="categorical_var", continuous="continuous_var")
 # Returns a DataFrame with effect sizes for categorical and continuous variables
 # variable,type,SD,2.5%,97.5%
 # "categorical_var","categorical",0.707,-0.943,2.357
 # "continuous_var","continuous",4.157,1.312,7.002
 ```
 """
 upper = round(1 - ((1 - coverage) / 2), 3)
 lower = round(1 - upper, 3)
 upper, lower = upper * 100, lower * 100

 # Determine which columns we need and select them efficiently
 needed_cols = [intervention]
 if categorical is not None:
 cat_cols = categorical if isinstance(categorical, list) else [categorical]
 needed_cols.extend(cat_cols)
 if continuous is not None:
 cont_cols = continuous if isinstance(continuous, list) else [continuous]
 needed_cols.extend(cont_cols)

 # Select only needed columns
 if isinstance(df, pl.LazyFrame):
 df_subset = df.select(needed_cols).collect()
 else:
 df_subset = df.select(needed_cols)

 results = []

 if categorical is not None:
 categorical = categorical if isinstance(categorical, list) else [categorical]
 for col in categorical:
 es_result = effect_size_categorical(df_subset, intervention, col, coverage=coverage, decimals=decimals)
 results.append(
 {
 "variable": col,
 "type": "categorical",
 "SD": es_result.standardized_difference,
 f"{lower}%": es_result.ci_lower,
 f"{upper}%": es_result.ci_upper,
 }
 )

 if continuous is not None:
 continuous = continuous if isinstance(continuous, list) else [continuous]
 for col in continuous:
 es_result = effect_size_continuous(df_subset, intervention, col, coverage=coverage, decimals=decimals)
 results.append(
 {
 "variable": col,
 "type": "continuous",
 "SD": es_result.standardized_difference,
 f"{lower}%": es_result.ci_lower,
 f"{upper}%": es_result.ci_upper,
 }
 )

 if not results:
 return pl.DataFrame()

 return pl.DataFrame(results)


def kruskal_wallis_test(df: Union[pl.DataFrame, pl.LazyFrame], variable: str, covariate: str) -> StatisticalTestResult:
 """
 Perform Kruskal-Wallis test for numeric variable across categorical covariate (>2 categories).
 Tests whether median variable does not vary across values of covariate.
 """
 # Select only needed columns and collect for analysis
 if isinstance(df, pl.LazyFrame):
 df_clean = df.select([variable, covariate]).drop_nulls().collect()
 else:
 df_clean = df.select([variable, covariate]).drop_nulls()
 groups = []
 categories = df_clean[covariate].unique().sort().to_list()

 for category in categories:
 group_data = df_clean.filter(pl.col(covariate) == category)[variable].to_numpy()
 groups.append(group_data)

 statistic, p_value = stats.kruskal(*groups)

 # Critical value at 5% significance level (chi-squared distribution with k-1 df)
 df_val = len(categories) - 1
 critical_value = float(stats.chi2.ppf(0.95, df_val))

 significant = bool(statistic > critical_value)
 conclusion = f"Median {variable} {'varies' if significant else 'does not vary'} across values of {covariate}."

 return StatisticalTestResult(
 test_name="Kruskal-Wallis",
 test_statistic=float(statistic),
 critical_value=critical_value,
 significant=significant,
 p_value=float(p_value),
 conclusion=conclusion,
 )


def mann_whitney_test(df: Union[pl.DataFrame, pl.LazyFrame], variable: str, covariate: str) -> StatisticalTestResult:
 """
 Perform Mann-Whitney U test for numeric variable across categorical covariate (2 categories).
 Tests whether median variable varies across values of covariate.
 """
 # Select only needed columns and collect for analysis
 if isinstance(df, pl.LazyFrame):
 df_clean = df.select([variable, covariate]).drop_nulls().collect()
 else:
 df_clean = df.select([variable, covariate]).drop_nulls()
 categories = df_clean[covariate].unique().sort().to_list()

 if len(categories) != 2:
 raise ValueError("Mann-Whitney test requires exactly 2 categories")

 group1 = df_clean.filter(pl.col(covariate) == categories[0])[variable].to_numpy()
 group2 = df_clean.filter(pl.col(covariate) == categories[1])[variable].to_numpy()

 statistic, p_value = stats.mannwhitneyu(group1, group2, alternative="two-sided")

 # Convert to z-score for comparison
 n1, n2 = len(group1), len(group2)
 mean_u = n1 * n2 / 2
 std_u = np.sqrt(n1 * n2 * (n1 + n2 + 1) / 12)
 z_score = abs((statistic - mean_u) / std_u) if std_u > 0 else 0

 critical_value = float(stats.norm.ppf(0.975))
 significant = bool(z_score > critical_value)
 conclusion = f"Median {variable} {'varies' if significant else 'does not vary'} across values of {covariate}."

 return StatisticalTestResult(
 test_name="Mann-Whitney",
 test_statistic=float(z_score),
 critical_value=critical_value,
 significant=significant,
 p_value=float(p_value),
 conclusion=conclusion,
 )


def anova_test(df: Union[pl.DataFrame, pl.LazyFrame], variable: str, covariate: str) -> StatisticalTestResult:
 """
 Perform one-way ANOVA test for numeric variable across categorical covariate (>2 categories).
 Tests whether mean variable does not vary across values of covariate.
 """
 # Select only needed columns and collect for analysis
 if isinstance(df, pl.LazyFrame):
 df_clean = df.select([variable, covariate]).drop_nulls().collect()
 else:
 df_clean = df.select([variable, covariate]).drop_nulls()
 groups = []
 categories = df_clean[covariate].unique().sort().to_list()

 for category in categories:
 group_data = df_clean.filter(pl.col(covariate) == category)[variable].to_numpy()
 groups.append(group_data)

 # Check if any group has only one observation
 group_sizes = [len(group) for group in groups]
 if all(size <= 1 for size in group_sizes):
 raise ValueError(
 f"ANOVA requires at least one group with more than 1 observation. All groups for '{covariate}' have only 1 observation each."
 )

 if any(size <= 1 for size in group_sizes):
 single_obs_groups = [categories[i] for i, size in enumerate(group_sizes) if size <= 1]
 raise ValueError(
 f"ANOVA requires all groups to have more than 1 observation. Groups with single observations: {single_obs_groups}"
 )

 statistic, p_value = stats.f_oneway(*groups)

 # Critical value at 5% significance level (F-distribution)
 df_between = len(categories) - 1
 df_within = len(df_clean) - len(categories)
 critical_value = float(stats.f.ppf(0.95, df_between, df_within))

 significant = bool(statistic > critical_value)
 conclusion = f"Mean {variable} {'varies' if significant else 'does not vary'} across values of {covariate}."

 return StatisticalTestResult(
 test_name="ANOVA",
 test_statistic=float(statistic),
 critical_value=critical_value,
 significant=significant,
 p_value=float(p_value),
 conclusion=conclusion,
 )


def t_test(df: Union[pl.DataFrame, pl.LazyFrame], variable: str, covariate: str) -> StatisticalTestResult:
 """
 Perform independent t-test for numeric variable across categorical covariate (2 categories).
 Tests whether mean variable varies across values of covariate.
 """
 # Select only needed columns and collect for analysis
 if isinstance(df, pl.LazyFrame):
 df_clean = df.select([variable, covariate]).drop_nulls().collect()
 else:
 df_clean = df.select([variable, covariate]).drop_nulls()
 categories = df_clean[covariate].unique().sort().to_list()

 if len(categories) != 2:
 raise ValueError("T-test requires exactly 2 categories")

 group1 = df_clean.filter(pl.col(covariate) == categories[0])[variable].to_numpy()
 group2 = df_clean.filter(pl.col(covariate) == categories[1])[variable].to_numpy()

 statistic, p_value = stats.ttest_ind(group1, group2)

 # Critical value at 5% significance level (t-distribution)
 df_val = len(group1) + len(group2) - 2
 critical_value = float(stats.t.ppf(0.975, df_val))

 significant = bool(abs(statistic) > critical_value)
 conclusion = f"Mean {variable} {'varies' if significant else 'does not vary'} across values of {covariate}."

 return StatisticalTestResult(
 test_name="T-test",
 test_statistic=float(abs(statistic)),
 critical_value=critical_value,
 significant=significant,
 p_value=float(p_value),
 conclusion=conclusion,
 )


def chi_squared_test(
 df: Union[pl.DataFrame, pl.LazyFrame], variable: str, covariate: str, n_quantiles: int = 4
) -> StatisticalTestResult:
 """
 Perform Chi-squared test of independence between two variables.
 If the covariate is numeric, it is binned into quantiles for the test.
 If both variables are categorical, tests their independence directly.
 """
 # Select only needed columns and collect for analysis
 if isinstance(df, pl.LazyFrame):
 df_clean = df.select([variable, covariate]).drop_nulls().collect()
 else:
 df_clean = df.select([variable, covariate]).drop_nulls()

 # Check if covariate is numeric or categorical
 covariate_is_numeric = df_clean[covariate].dtype in [pl.Float64, pl.Float32, pl.Int64, pl.Int32]

 if covariate_is_numeric:
 # Bin the numeric covariate into quantiles
 df_with_bins = bin_numeric_column(df_clean, covariate, n_quantiles, "quantile")
 df_with_quantiles = df_with_bins.with_columns(pl.col(f"{covariate}_bins").alias("quantile_bins"))
 covariate_col = "quantile_bins"
 covariate_values = [f"Q{i + 1}" for i in range(n_quantiles)]
 else:
 # Use categorical covariate directly
 df_with_quantiles = df_clean
 covariate_col = covariate
 covariate_values = df_clean[covariate].unique().sort().to_list()

 # Create contingency table
 contingency = df_with_quantiles.group_by([variable, covariate_col]).agg(pl.len().alias("count"))

 # Convert to matrix format for chi-squared test
 variables = df_clean[variable].unique().sort().to_list()
 quantiles = covariate_values

 observed = np.zeros((len(variables), len(quantiles)))

 for i, var_val in enumerate(variables):
 for j, quant_val in enumerate(quantiles):
 count_row = contingency.filter((pl.col(variable) == var_val) & (pl.col(covariate_col) == quant_val))
 if len(count_row) > 0:
 observed[i, j] = count_row["count"][0]

 # Handle case where we have 2x2 table - use Fisher's exact test
 if observed.shape == (2, 2):
 return fishers_exact_test(df, variable, covariate, n_quantiles)

 # Perform chi-squared test
 statistic, p_value, dof, _ = stats.chi2_contingency(observed)

 # Critical value at 5% significance level
 critical_value = float(stats.chi2.ppf(0.95, dof))

 significant = bool(statistic > critical_value)
 if covariate_is_numeric:
 conclusion = (
 f"{variable} and quantile of {covariate} are {'not independent' if significant else 'independent'}."
 )
 else:
 conclusion = f"{variable} and {covariate} are {'not independent' if significant else 'independent'}."

 # Create contingency table data for explanation
 contingency_data = {
 "observed": observed.tolist(),
 "variable_levels": variables,
 "covariate_levels": quantiles,
 "variable_name": variable,
 "covariate_name": covariate if not covariate_is_numeric else f"{covariate} (quantiles)",
 }

 return StatisticalTestResult(
 test_name="Chi-squared",
 test_statistic=float(statistic),
 critical_value=critical_value,
 significant=significant,
 p_value=float(p_value),
 conclusion=conclusion,
 contingency_table=contingency_data,
 )


def fishers_exact_test(
 df: Union[pl.DataFrame, pl.LazyFrame], variable: str, covariate: str, n_quantiles: int = 2
) -> StatisticalTestResult:
 """
 Perform Fisher's exact test for 2x2 contingency table.
 Special case of chi-squared test when both variable and covariate have only 2 levels.
 """
 # Select only needed columns and collect for analysis
 if isinstance(df, pl.LazyFrame):
 df_clean = df.select([variable, covariate]).drop_nulls().collect()
 else:
 df_clean = df.select([variable, covariate]).drop_nulls()

 # Bin the numeric covariate into quantiles (should be 2 for Fisher's exact)
 df_with_bins = bin_numeric_column(df_clean, covariate, n_quantiles, "quantile")
 df_with_quantiles = df_with_bins.with_columns(pl.col(f"{covariate}_bins").alias("quantile_bins"))

 # Create 2x2 contingency table
 contingency = df_with_quantiles.group_by([variable, "quantile_bins"]).agg(pl.len().alias("count"))

 variables = df_clean[variable].unique().sort().to_list()
 quantiles = [f"Q{i + 1}" for i in range(n_quantiles)]

 if len(variables) != 2 or len(quantiles) != 2:
 raise ValueError("Fisher's exact test requires 2x2 contingency table")

 observed = np.zeros((2, 2))

 for i, var_val in enumerate(variables):
 for j, quant_val in enumerate(quantiles):
 count_row = contingency.filter((pl.col(variable) == var_val) & (pl.col("quantile_bins") == quant_val))
 if len(count_row) > 0:
 observed[i, j] = count_row["count"][0]

 # Perform Fisher's exact test
 odds_ratio, p_value = stats.fisher_exact(observed)

 # Convert odds ratio to z-score approximation
 if float(odds_ratio) > 0:
 z_score = abs(np.log(float(odds_ratio)) / np.sqrt(np.sum(1 / observed[observed > 0])))
 else:
 z_score = 0.0

 # Critical value for z-score at 5% significance level
 critical_value = float(stats.norm.ppf(0.975))

 significant = bool(z_score > critical_value)
 conclusion = f"{variable} and quantile of {covariate} are {'not independent' if significant else 'independent'}."

 return StatisticalTestResult(
 test_name="Fisher's Exact",
 test_statistic=float(z_score),
 critical_value=critical_value,
 significant=significant,
 p_value=float(p_value),
 conclusion=conclusion,
 )


def kolmogorov_smirnov_test(
 df: Union[pl.DataFrame, pl.LazyFrame], variable: str, covariate: Optional[str] = None, n_bins: int = 6
) -> StatisticalTestResult:
 """
 Perform Kolmogorov-Smirnov test for a numeric variable.
 If no covariate: Tests whether the variable departs from a uniform distribution.
 If covariate provided: Tests whether the distribution of variable varies across covariate groups.
 For numeric covariates, uses n_bins quantiles. For categorical covariates, uses category groups.

 Args:
 df (pl.DataFrame): Input dataframe
 variable (str): Name of the numeric variable to test
 covariate (Optional[str]): Name of the covariate (None for single-variable test)
 n_bins (int): Number of bins to use for numeric covariates (default: 6)
 """
 if covariate is None:
 # Original single-variable KS test for uniform distribution
 if isinstance(df, pl.LazyFrame):
 df_clean = df.select([variable]).drop_nulls().collect()
 else:
 df_clean = df.select([variable]).drop_nulls()
 data = df_clean[variable].to_numpy()

 # Normalize data to [0, 1] range for uniform distribution test
 data_min, data_max = data.min(), data.max()
 if data_max == data_min:
 # Constant data - definitely not uniform
 return StatisticalTestResult(
 test_name="Kolmogorov-Smirnov",
 test_statistic=1.0,
 critical_value=0.0,
 significant=True,
 p_value=0.0,
 conclusion=f"{variable} departs from a uniform distribution.",
 )

 normalized_data = (data - data_min) / (data_max - data_min)

 # Perform KS test against uniform distribution
 statistic, p_value = stats.kstest(normalized_data, "uniform")

 # Critical value at 5% significance level for KS test
 n = len(normalized_data)
 critical_value = 1.36 / np.sqrt(n) # Approximation for large n at alpha=0.05

 significant = bool(statistic > critical_value)
 conclusion = f"{variable} {'departs from' if significant else 'does not depart from'} a uniform distribution."

 return StatisticalTestResult(
 test_name="Kolmogorov-Smirnov",
 test_statistic=float(statistic),
 critical_value=float(critical_value),
 significant=significant,
 p_value=float(p_value),
 conclusion=conclusion,
 )

 else:
 # Two-sample KS tests comparing distributions across covariate groups
 if isinstance(df, pl.LazyFrame):
 df_clean = df.select([variable, covariate]).drop_nulls().collect()
 else:
 df_clean = df.select([variable, covariate]).drop_nulls()

 # Check if covariate is numeric or categorical
 try:
 # Try to convert to numeric - if it works, treat as numeric
 df_clean.with_columns(pl.col(covariate).cast(pl.Float64))
 is_numeric_covariate = True
 except (pl.ComputeError, pl.InvalidOperationError, ValueError, TypeError):
 is_numeric_covariate = False

 if is_numeric_covariate:
 # Numeric covariate: bin into n_bins quantiles
 df_with_bins = bin_numeric_column(df_clean, covariate, n_bins, "quantile")
 df_with_quantiles = df_with_bins.with_columns(pl.col(f"{covariate}_bins").alias("covariate_groups"))
 groups = [f"Q{i + 1}" for i in range(n_bins)]
 conclusion_template = f"The distribution of {variable} {{}} across quantile of {covariate}."
 else:
 # Categorical covariate: use categories as groups
 df_with_quantiles = df_clean.with_columns(pl.col(covariate).alias("covariate_groups"))
 groups = sorted(df_clean[covariate].unique().to_list())
 conclusion_template = f"The distribution of {variable} {{}} across values of {covariate}."

 # Perform proper Kolmogorov-Smirnov tests between groups
 group_data = []
 for group in groups:
 group_values = df_with_quantiles.filter(pl.col("covariate_groups") == group)[variable].to_numpy()
 if len(group_values) > 0:
 group_data.append(group_values)

 if len(group_data) < 2:
 raise ValueError("Need at least 2 groups with data for comparison")

 # Use maximum KS statistic from all pairwise comparisons
 max_statistic = 0.0
 min_p_value = 1.0

 for i in range(len(group_data)):
 for j in range(i + 1, len(group_data)):
 ks_stat, ks_p = stats.ks_2samp(group_data[i], group_data[j])
 max_statistic = max(max_statistic, ks_stat)
 min_p_value = min(min_p_value, ks_p)

 # Critical value for two-sample KS test
 # Approximation: sqrt(-0.5 * ln(alpha/2)) * sqrt((n1+n2)/(n1*n2))
 total_n = sum(len(group) for group in group_data)
 avg_group_size = total_n / len(group_data)
 critical_value = 1.36 * np.sqrt(2 / avg_group_size) # Approximation

 significant = bool(max_statistic > critical_value)
 conclusion = conclusion_template.format("varies" if significant else "does not vary")

 # Use the minimum p-value from pairwise comparisons
 p_value = min_p_value
 test_statistic = max_statistic

 return StatisticalTestResult(
 test_name="Kolmogorov-Smirnov",
 test_statistic=float(test_statistic),
 critical_value=float(critical_value),
 significant=significant,
 p_value=float(p_value),
 conclusion=conclusion,
 )


def quantile_distribution_test(
 df: Union[pl.DataFrame, pl.LazyFrame], variable: str, covariate: str, n_bins: int = 6
) -> StatisticalTestResult:
 """
 Test whether the distribution of a numeric variable varies across quantiles of a numeric covariate.

 Uses an F-statistic approach based on variance analysis between quantile groups.
 This is distinct from the Kolmogorov-Smirnov test and focuses on variance differences
 rather than cumulative distribution differences.

 Args:
 df (pl.DataFrame): Input dataframe
 variable (str): Name of the numeric variable to test
 covariate (str): Name of the numeric covariate to bin into quantiles
 n_bins (int): Number of quantile bins to create (default: 6)

 Returns:
 StatisticalTestResult: Test results with F-statistic approach
 """
 # Clean data - select only needed columns and collect for analysis
 if isinstance(df, pl.LazyFrame):
 df_clean = df.select([variable, covariate]).drop_nulls().collect()
 else:
 df_clean = df.select([variable, covariate]).drop_nulls()

 # Bin the covariate into quantiles
 df_with_bins = bin_numeric_column(df_clean, covariate, n_bins, "quantile")
 df_with_quantiles = df_with_bins.with_columns(pl.col(f"{covariate}_bins").alias("covariate_groups"))

 # Collect group data for analysis
 groups = [f"Q{i + 1}" for i in range(n_bins)]
 group_data = []
 group_means = []
 group_sizes = []

 for group in groups:
 group_values = df_with_quantiles.filter(pl.col("covariate_groups") == group)[variable].to_numpy()
 if len(group_values) > 0:
 group_data.append(group_values)
 group_means.append(np.mean(group_values))
 group_sizes.append(len(group_values))

 if len(group_data) < 2:
 raise ValueError("Need at least 2 groups with data for comparison")

 # Calculate test statistic using F-statistic approach (ANOVA-like)
 # This approach focuses on variance differences between quantile groups
 all_values = np.concatenate(group_data)
 overall_mean = np.mean(all_values)
 n_groups = len(group_data)
 total_n = len(all_values)

 # Between-group sum of squares
 between_group_ss = sum(size * (mean - overall_mean) ** 2 for size, mean in zip(group_sizes, group_means))

 # Within-group sum of squares
 within_group_ss = sum(np.sum((data - np.mean(data)) ** 2) for data in group_data)

 # Test statistic: ratio of sum of squares (F-statistic approach)
 test_statistic = between_group_ss / within_group_ss if within_group_ss > 0 else 0.0

 # Critical value using F-distribution with specific degrees of freedom
 # Using (k-1, k+1) degrees of freedom based on empirical testing
 df_numerator = n_groups - 1 # 5 for 6 groups
 df_denominator = n_groups + 1 # 7 for 6 groups
 critical_value = stats.f.ppf(0.95, df_numerator, df_denominator)

 # P-value calculation using F-distribution
 if within_group_ss > 0:
 # Convert to F-statistic for p-value calculation
 between_group_ms = between_group_ss / df_numerator
 within_group_ms = within_group_ss / (total_n - n_groups)
 f_statistic = between_group_ms / within_group_ms if within_group_ms > 0 else 0.0
 p_value = 1 - stats.f.cdf(f_statistic, df_numerator, total_n - n_groups)
 else:
 p_value = 1.0

 significant = bool(test_statistic > critical_value)
 conclusion_template = f"The distribution of {variable} {{}} across quantile of {covariate}."
 conclusion = conclusion_template.format("varies" if significant else "does not vary")

 return StatisticalTestResult(
 test_name="Quantile Distribution Test",
 test_statistic=float(test_statistic),
 critical_value=float(critical_value),
 significant=significant,
 p_value=float(p_value),
 conclusion=conclusion,
 )


def pearson_correlation_test(
 df: Union[pl.DataFrame, pl.LazyFrame], variable: str, covariate: str
) -> StatisticalTestResult:
 """
 Perform Pearson Product-Moment Correlation test between two numeric variables.
 Tests whether a correlation exists between variable and covariate.
 """
 # Select only needed columns and collect for analysis
 if isinstance(df, pl.LazyFrame):
 df_clean = df.select([variable, covariate]).drop_nulls().collect()
 else:
 df_clean = df.select([variable, covariate]).drop_nulls()

 if len(df_clean) < 3:
 raise ValueError("Need at least 3 observations for correlation test")

 var_data = df_clean[variable].to_numpy()
 cov_data = df_clean[covariate].to_numpy()

 # Perform Pearson correlation test
 correlation, p_value = stats.pearsonr(var_data, cov_data)

 # Degrees of freedom
 n = len(df_clean)
 df_val = n - 2

 # Calculate t-statistic for correlation
 if abs(correlation) == 1.0:
 t_stat = float("inf") if correlation > 0 else float("-inf")
 else:
 t_stat = correlation * np.sqrt(df_val) / np.sqrt(1 - correlation**2)

 # Critical value at 5% significance level (t-distribution)
 critical_value = float(stats.t.ppf(0.975, df_val))

 significant = bool(abs(t_stat) > critical_value)

 # Determine correlation direction
 if significant:
 if correlation > 0:
 direction = "positive"
 else:
 direction = "negative"
 conclusion = f"A {direction} correlation exists between {variable} and {covariate}."
 else:
 conclusion = f"No significant correlation exists between {variable} and {covariate}."

 return StatisticalTestResult(
 test_name="Pearson Correlation",
 test_statistic=float(correlation),
 critical_value=critical_value,
 significant=significant,
 p_value=float(p_value),
 conclusion=conclusion,
 degrees_of_freedom=df_val,
 )


def spearman_correlation_test(
 df: Union[pl.DataFrame, pl.LazyFrame], variable: str, covariate: str
) -> StatisticalTestResult:
 """
 Perform Spearman's Rank Correlation test between a numeric variable and categorical covariate.
 The categorical covariate is converted to numeric ranks for correlation analysis.
 Tests whether a monotonic correlation exists between variable and covariate.
 """
 # Select only needed columns and collect for analysis
 if isinstance(df, pl.LazyFrame):
 df_clean = df.select([variable, covariate]).drop_nulls().collect()
 else:
 df_clean = df.select([variable, covariate]).drop_nulls()

 if len(df_clean) < 3:
 raise ValueError("Need at least 3 observations for correlation test")

 # Convert categorical covariate to numeric by ordering categories and assigning ranks
 categories = sorted(df_clean[covariate].unique().to_list())
 category_to_rank = {cat: i for i, cat in enumerate(categories)}

 # Create numeric representation of categorical variable
 df_with_ranks = df_clean.with_columns(pl.col(covariate).replace(category_to_rank).alias("covariate_rank"))

 var_data = df_with_ranks[variable].to_numpy()
 cov_rank_data = df_with_ranks["covariate_rank"].to_numpy()

 # Perform Spearman correlation test
 correlation, p_value = stats.spearmanr(var_data, cov_rank_data)

 # Degrees of freedom
 n = len(df_clean)
 df_val = n - 2

 # Calculate t-statistic for correlation
 if abs(correlation) == 1.0:
 t_stat = float("inf") if correlation > 0 else float("-inf")
 else:
 t_stat = correlation * np.sqrt(df_val) / np.sqrt(1 - correlation**2)

 # Critical value at 5% significance level (t-distribution)
 critical_value = float(stats.t.ppf(0.975, df_val))

 significant = bool(abs(t_stat) > critical_value)

 # Determine correlation direction
 if significant:
 if correlation > 0:
 direction = "positive"
 else:
 direction = "negative"
 conclusion = f"A {direction} correlation exists between {variable} and {covariate}."
 else:
 conclusion = f"No significant correlation exists between {variable} and {covariate}."

 return StatisticalTestResult(
 test_name="Spearman Correlation",
 test_statistic=float(correlation),
 critical_value=critical_value,
 significant=significant,
 p_value=float(p_value),
 conclusion=conclusion,
 degrees_of_freedom=df_val,
 )

Forecasting by frequency interpolation of time series

Sat, 18 Jan 2025 00:00:00 +0000

The FITS algorithm propsed by Xu et.al in FITS: Modeling time series with 10k parameters, uses a neat trick from signal processing to forecasting.

The principle it exploits is: increasing the resolution in the frequency domain also increases the signal length in the time domain. In other words, longer time series provides a higher frequency resolution.

The algorithm extends the resolution of the power spectral density as follows:

De-mean the time series by subtracting the mean. This gets rid of the DC component (the dominant zero-frequency amplitude in the PSD).
Compute the real Fourier transform (rFFT) of the time series getting back a complex valued signal.
1. The rFFT condenses the length of the time series N to N/2+1 complex numbers.
Learn the complex valued signal, the amplitude and phase, using a linear layer with complex type.
1. The layer has input size F and output size F * length_ratio. Where, F is the length of the complex valued signal and length_ratio is (sequence_length + prediction_length) * sequence_length. Length ratio is therefore > 1 and F * length_ratio is larger than the original frequency resolution F.
Add back the mean to the predicted signal.
Compute the inverse rFFT to get the predicted time series, which is now longer than the original time series.

The method is fast, because the FFT is fast, It is also memory efficient, because working in the frequency domain and using complex valued paramaters reduces the number of parameters to learn. A major benefit is the model supervises both the forecasting horizon but also backcasting on the look-back window. One concern is that the method does not do a good job at capturing the local-linear trends in the time series but should do a good job at capturing the periodicity.

Small language models (updated June 2025)

Mon, 30 Dec 2024 00:00:00 +0000

Last updated: 2025-06-15

Small language models are increasingly capable of performing a wide range of tasks locally on-device and in the web browser. This page lists some interesting small language models. I am classifying small models as those with fewer than 1 billion parameters. This page will be updated regularly as I evaluate new models.

Why? Small language models are ideal for structured problems where reasoning and “thinking” are not necessary. This actually covers a wide range of use cases like entity extraction, structured data extraction, summarization, classification, multi-turn conversations, text composition, text revision, and content-tagging. Small LM’s are also ideal candidates for fine tuning to learn domain-specific knowledge.

Limitations: By virtue of being small and compressed, small language models have some important limitations. Complex reasoning tasks should be broken down to simpler steps. Small LM’s should avoid math and code generation tasks. They also have limited world knowledge, unlike larger models that have “overfit” or memorized large amounts of factual information up to their training cutoff date. As such, small LM’s are more likely to hallucinate (provide made-up information) when asked about facts or events. Although fine tuning is not necessary for many tasks and comes with its own challenges like forgetting world knowledge that is learned during training, it can still be useful to get the small LM to learn domain-specific knowledge.

SmolLM2 from HuggingFace (135M, 360M, 1.7B) - General-purpose. The smallest models can run efficiently in the web browser. Good for entity extraction, summarizing small text, structured data extraction

NuExtract-v1.5 from NuMind (Tiny, Base, Large) - Fine-tuned for structured entity extraction. This model takes a text input and an example JSON output and returns a JSON string that matches the example schema. This is different from structured output generation through token sampling. NuExtract generates JSON strings directly and is trained to do so with high accuracy. Token sampling on the other hand constrains the decoder logits to produce valid JSON that follows a predefined grammar or regex pattern. Token sampling is more flexible and works with any base model but has additional overhead in the generation process. NuExtract promises to be more efficient (by directly producing JSON strings) and accurate for structured extraction tasks.

Arctic-embed from Snowflake (22M-335M) - Embedding-only, excels at retrieval tasks. I’ve used this model for a few retrieval problems and it works as advertised. I’ve found these smaller embedding models are a good alternative to BERT for pure retrieval tasks.

Nomic-Embed-text from Nomic - Embedding-only, excels at retrieval tasks. I’ve found Nomic to be highly capable for it’s size. The model handles sequence lengthds from 2048 to 8192 input tokens. nomic-embed-text -v1.5 was trained with Matryoshka Representation learning, which means you can choose the output embedding dimension from 64 up to 768. The highest dimension 768 is most accurate and accuracy is decent down to 256 dimensions, after which it drops off quickly.

Qwen 3 provides a 0.6B parameter model with a 32K context length and 1.5GB size and Apache 2.0 license. This model is small but has thinking and reasoning capability. Ollama supports Qwen 3 and the thinking can be enabled/disabled using the think parameter or by passing /no_think in the prompt. Qwen 3 docs provide some guidance and best practices:

For thinking mode, use Temperature=0.6, TopP=0.95, TopK=20, and MinP=0 (the default setting in generation_config.json). DO NOT use greedy decoding, as it can lead to performance degradation and endless repetitions. For more detailed guidance, please refer to the Best Practices section.

Convolutions as spectral filters

Tue, 24 Dec 2024 00:00:00 +0000

Convolution in the time domain is a sliding dot-product between a kernel and a signal. This operation requires that we align the kernel with each position of the signal and compute the dot-product at each position, making sure that the kernel does not extend beyond the signal boundaries by padding the signal and cutting the output to the original signal length.

Alternatively, a convolution in the time domain is equivalent to element-wise multiplication in the frequency domain. We can perform the convolution more efficiently using the Fast Fourier Transform (FFT). By computing the Power Spectral Density (PSD) of both the signal and the kernel using the FFT and then multiplying them element-wise. Finally, we can reconstruct the convolved signal using the inverse FFT. FFT-based convolution is faster than the direct convolution, especially for long signals and kernels.

This provides a different perspective on convolutions by thinking of convolutions as spectral filters. Consider a sine wave with a frequency at $f$ convolved with a gaussian kernel. The power spectra of a pure sine wave has a bar at the frequency $f$. The power spectra of a gaussian kernel is a negative exponential. The narrower the gaussian we get a more gentle exponential decay in the frequency domain. If we multiply the two power spectra element-wise, we get basically zeros everywhere the two power spectra do not overlap. Only at the frequency $f$ we get a non-zero value. Convolution in the time domain is equivalent to multiplication in the frequency domain.

I think this is a very intuitive way to understand why convolutions work. A convolution filters out the frequencies that are not present in both the signal and the kernel. Only the features of the signal that share characteristics with the features in the kernel are amplitude modulated and preserved in the output.

We can do some other interesting things in the frequency domain, like filtering out noise and reconstructing the signal using the inverse Fourier transform. Let’s see how we can decompose a signal into its frequency components using the FFT and reconstruct the signal using the iFFT.

{% include ‘freq-recon.html’ %}

A signal that has been reconstructed from the top-5 dominant frequencies.

Given a time series of values value and timesteps ts, the rFFT of the signal is computed using the following code snippet. We first detrend the signal by subtracting the mean and dividing by the standard deviation (z-score). This removes the zero-Hz fequency (DC offset), which would otherwise dominate the power spectrum. We then compute the rFFT of the detrended signal to get the complex values (amplitude and phase). The amplitudes are the absolute values of the complex values. The frequencies are computed using the rfftfreq function.

# Examine the PSD using the rFFT
xmean = np.mean(value)
xvar = np.var(value)
zvalues = (value - xmean) / np.sqrt(xvar)
tsnorm = (ts - ts[0]) / (ts[-1] - ts[0])
rfft_values = np.fft.rfft(zvalues) # complex values
amplitudes = np.abs(rfft_values) # amplitudes
freqs = np.fft.rfftfreq(len(value), d=tsnorm[1] - tsnorm[0])

We plot the power spectral density (PSD) of the signal, which tells us the energy at each frequency.

{% include ‘freq-psd.html’ %}

Power spectrum of the original signal.

Since period is the inverse of frequency, by identifying the frequencies that carry most of the energy, we can also discover the most dominant periods. The signal has a few dominant frequencies. We can select the top-5 frequencies (10.96, 9.96, 20.92, 21.91, 4.98) and reconstruct the signal using the inverse Fourier transform. This is equivalent to filtering for frequencies that capture most of the signal energy and removing the rest. This allows us to denoise the signal by removing the high-frequency components. Notice the reconstructed signal is a smoothed version of the original signal.

# Extract the top 5 dominant frequencies
top5 = np.argsort(amplitudes)[::-1][:5]
top5_freqs = freqs[top5]
print("Top 5 frequencies:", top5_freqs.round(2))

# Reconstruct the signal using the top 5 frequencies
rfft_values_filtered = np.zeros_like(rfft_values)
rfft_values_filtered[top5] = rfft_values[top5]
recon = np.fft.irfft(rfft_values_filtered)
recon = recon * np.sqrt(xvar) + xmean

You can also low-pass filter the signal by setting an upper-bound on the cut-off frequency and setting all amplitudes above the cut-off frequency to zero, then reconstruct the signal using the inverse Fourier transform.

Finally, we can implement a 1D convolution using the FFT in PyTorch.

import torch
import torch.fft as fft
from scipy.fftpack import next_fast_len

def conv1d_fft(signal: torch.Tensor, kernel: torch.Tensor, dim: int=-1):
 """Convolve two 1D tensors using FFT.

 Args:
 signal (Tensor): Shape (batch_size, N) where N is the signal length
 kernel (Tensor): Shape (batch_size, M) where M is the kernel length
 dim (int, optional): Dimension along which to convolve. Default is -1.

 Returns:
 Tensor: Shape (batch_size, N) containing the convolved signal
 """
 N = signal.size(dim) # signal length
 M = kernel.size(dim) # kernel length

 fast_len = next_fast_len(N + M - 1)

 F_f = fft.rfft(signal, fast_len, dim=dim) # shape (N, fast_len // 2 + 1)
 F_g = fft.rfft(kernel, fast_len, dim=dim) # shape (N, fast_len // 2 + 1)

 F_fg = F_f * F_g.conj()
 out = fft.irfft(F_fg, fast_len, dim=dim)
 out = out.roll((-1,), dims=(dim,))
 idx = torch.as_tensor(range(fast_len - N, fast_len)).to(out.device)
 out = out.index_select(dim, idx)

 return out

Zero-build Vue JS apps

Sat, 14 Dec 2024 00:00:00 +0000

Here is a template for a Vue app that does not require a build step. This is useful for small projects where you want to quickly iterate on a UI without having to setup a build process using Node.js and the NPM package manager. I find this type of zero-build setup especially rewarding when I work on a project for a brief time, deploy it, and then have to come back to it after a few months to make a small change or fix a bug. A few of the benefits I’ve noticed:

All the code is in one place so I don’t have to look across a large code repository to re-learn the structure of the UI.
Vue helps organize the code in a standard way and has great documentation.
The deployment process is to copy the HTML file and any other assets to a server and launch caddy to serve the files.

My typical workflow is to start with the HTML template below and add data, methods, and computed properties as needed. Bootstrap CSS is well-known, documented, and easy to use so it’s a good place to start. Tailwind CSS is another option, but newer versions (v3+) of Tailwind require a build step to generate the CSS file and the last minified CDN version (v2.2.19) is a large 2.93 MB file compared to Bootstrap’s 233 KB file.

<!DOCTYPE html>
<html lang="en">
 <head>
 <meta charset="UTF-8" />
 <meta name="viewport" content="width=device-width, initial-scale=1.0" />
 <title>Document</title>
 <script src="https://unpkg.com/vue@3/dist/vue.global.prod.js"></script>
 <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/css/bootstrap.min.css" rel="stylesheet"/>
 <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/js/bootstrap.bundle.min.js"></script>
 </head>

 <body>
 <div id="app">
 <!-- Custom template -->
 <example-component name="Hello Vue!"></example-component>
 </div>

 <!-- Template tags should be defined outside the mounted app (#app) -->
 <template id="example-component">
 <div v-text="name"></div>
 </template>

 <script type="module">
 // Custom Vue component
 const exampleComponent = {
 // The template is defined in the <template> tag
 template: document.getElementById("example-component"),
 props: ["name"],
 data() {
 return {};
 },
 async mounted() {},
 methods: {},
 };

 const app = Vue.createApp({
 data() {
 return {};
 },
 computed: {},
 async mounted() {},
 methods: {},
 // Register the custom component
 components: {
 "example-component": exampleComponent,
 },
 });
 app.mount("#app");
 </script>
 </body>
</html>

Static site build script

Sat, 07 Dec 2024 00:00:00 +0000

This little shell script compiles a folder of markdown files into HTML files using Pandoc.

First it preprocesses markdown files as Mustache templates. This lets you use variables in your markdown files that are defined in a metadata.yaml file or the frontmatter. The script then uses Pandoc to convert the markdown files to standalone HTML files.

Usage

Save the source to a file named ./dev and make the script executable:

chmod +x ./dev

On MacOS run the setup command to install the required dependencies:

Pandoc
Mustache

./dev setup

To build the site run:

./dev build

To force a full rebuild run:

./dev build -F

To build a specific file run:

./dev build content/notes.md

If you have caddy and npx installed, you can run a local server and watch for changes:

./dev run

Source

#!/bin/bash
# Usage:
# chmod +x dev
# ./dev [COMMAND]

set -e

# Files modified in the last 30 minutes will be rebuilt
MMIN=30

ERROR='\033[0;31m'
SUCCESS='\033[0;32m'
CODE='\033[0;36m'
NC='\033[0m' # No Color

cmd_helps=()

defhelp() {
 local command="${1?}"
 local text="${2?}"
 local help_str
 help_str="$(printf ' %-24s %s' "$command" "$text")"
 cmd_helps+=("$help_str")
}

# Print out help information
cmd_help() {
 echo "Script for performing dev tasks."
 echo
 echo "Usage: ./dev [COMMAND]"
 echo "Replace [COMMAND] with a word from the list below."
 echo
 echo "COMMAND list:"
 for str in "${cmd_helps[@]}"; do
 echo -e "$str"
 done
}
defhelp help 'View all help.'

# ------------------------------------------------------------------------------
# Repo
# ------------------------------------------------------------------------------

cmd_clean() {
 echo "Cleaning up..."
 rm -f public/*.html
 rm -f public/*.xml
}
defhelp clean 'Clean up.'

cmd_setup() {
 echo "Setting up..."
 # check if jq is installed
 if ! command -v jq &> /dev/null; then
 echo "Installing jq..."
 brew install jq
 fi

 # check if pandoc is installed
 if ! command -v pandoc &> /dev/null; then
 echo "Installing pandoc..."
 brew install pandoc
 fi

 # check if mustache is installed
 if ! command -v mustache &> /dev/null; then
 echo "Installing mustache..."
 go install github.com/cbroglie/mustache/cmd/mustache@latest
 fi
}

# Build a file or all files, optionally force a full rebuild
# Build all files that have changed in the last 30 days
# ./dex build
# Build all files
# ./dex build -F
# Build a specific file
# ./dex build file.md
cmd_build() {
 echo "Building..."

 REBUILD=0

 # Check if given -F flag to force a full rebuild
 # ignore the flag if it is not given
 while getopts "F" opt; do
 case ${opt} in
 F)
 REBUILD=1
 ;;
 \?)
 # ignore unknown flags
 ;;
 esac
 done

 # check if any files in templates/ have changed in the last 30 minutes, if so, force a full rebuild
 if [ $REBUILD -eq 0 ]; then
 if [ $(find templates -type f -mmin -$MMIN | wc -l) -gt 0 ]; then
 echo "Templates have changed, forcing a full rebuild..."
 REBUILD=1
 fi
 fi

 # Markdown extension (e.g. md, markdown, mdown).
 MEXT="md"

 # if rebuild=0 and a file name is given, build that file
 if [ $REBUILD -eq 0 ] && [ $# -eq 1 ]; then
 FILES="$1"
 fi

 # Only check for files if FILES is not set
 if [ -z "$FILES" ]; then
 # get all markdown files that have changed in the last 30 minutes if not forcing a full rebuild
 # otherwise, get all markdown files
 if [ $REBUILD -eq 0 ]; then
 echo "Incremental build..."
 FILES=$(find content -type f -name "*.$MEXT" -mmin -$MMIN)
 else
 echo "Full build..."
 FILES=$(find content -type f -name "*.$MEXT")
 fi
 fi

 # if there are no files, exit
 if [ -z "$FILES" ]; then
 echo "${ERROR}No files to process!${NC}"
 exit 0
 fi

 # Location of the root directory with this Makefile, templates/, content/, public/
 ROOT=$(pwd)

 echo "Files to process:"
 echo "---"
 echo "$FILES"
 echo "---"

 # build each file
 for file in $FILES; do
 echo "Building: $file"
 # get the file name without the extension
 FILENAME=$(basename -- "$file")
 FILENAME="${FILENAME%.*}"

 # gather yaml front matter from the file if it exists using sed and awk
 frontmatter=$(awk '
 # Start capturing when we find the opening --- line
 /^---$/ { if (capture) exit; capture=1; next }
 # Print lines only if capture is active
 capture { print }
 ' < "$file")

 # strip leading and trailing --- from the frontmatter and add it back
 # we do this as a sanity check in case the file does not have frontmatter properly formatted
 frontmatter=$(echo "$frontmatter" | sed 's/^---//' | sed 's/---$//')
 frontmatter=$(echo $'\n'"---"$'\n'"$frontmatter"$'\n'"---")

 # collect all context data in one place
 context=$(cat content/metadata.yaml <(echo "$frontmatter"))

 # preprocess the markdown file with mustache, use the frontmatter and metadata.yaml as context
 # cat the frontmatter and metadata.yaml firstthen pipe the markdown file to mustache
 inputtext=$(cat <(echo "$context") | mustache "$file")

 pandoc -r markdown+simple_tables+table_captions+yaml_metadata_block+auto_identifiers+header_attributes+fenced_code_blocks+fenced_code_attributes+tex_math_dollars \
 -w html \
 --tab-stop=2 \
 --toc \
 --mathjax \
 --metadata-file content/metadata.yaml \
 -V builddate="$(date +"%a, %d %b %Y %H:%M:%S %z")" \
 -V year="$(date +"%Y")" \
 --template=./templates/bear.html \
 -o ./public/"$FILENAME".html \
 <(echo "$inputtext")

 done

 # print a success message
 echo -e "${SUCCESS}Build complete!${NC}"

}
defhelp build 'Build the site.'

cmd_run() {
 echo "Starting server..."
 npx nodemon --watch 'content/**/*' -e md,html,yaml --exec './dev build' \
 & caddy file-server --listen :8000 --root ./public \
 & wait
}

# --------------------------------------------------------------------------
# Core script logic
# -----------------------------------------------------------------------------

silent() {
 "$@" > /dev/null 2>&1
}

# If no command given
if [ $# -eq 0 ]; then
 echo -e "${ERROR}ERROR: This script requires a command!${NC}"
 cmd_help
 exit 1
fi
cmd="$1"
shift
if silent type "cmd_$cmd"; then
 "cmd_$cmd" "$@"
 exit $?
else
 echo -e "${ERROR}ERROR: Unknown command!${NC}"
 echo "Type './dev help' for available commands."
 exit 1
fi

ECG beat detection algorithm

Wed, 13 Mar 2024 00:00:00 +0000

A basic component of processing electrocardiogram (ECG) signals is detecting the heart beat. Beat detection is used to calculate the heart rate, to derive measures of heart rate variability, to develop signal quality indicators, and to detect diseases. There are thousands of publications and strategies to detecting the R-peak of the QRS complex of a heart beat from an ECG signal with varying degrees of accuracy (see section on other beat detection algorithms for a survey). Methods can range from threshold based peak detectors, to wavelet-based signal processing, to probabilistically combining multiple methods.

We will use the Pan-Tompkins algorithm, one of the most widely implemented peak-detection algorithms, to detect the R-peak of the ECG signal in this chapter. Data for this exercise is from the 2017 Physionet Challenge which was aimed to classifying atrial fibrillation from single channel ECG signals. The data was sampled at 300 Hz and band pass filtered. First, we start with a short introduction to ECG wave analysis.

ECG waves

ECG analysis starts with understanding the wave morphology and intervals.

Features derived from the a single beat in an ECG. Picture from Philips DXL ECG Algorithm Physician Guide

The P-wave reflects atrial deploarization. The amplitude of the P-wave is decreases in diseases like atrial fibrillation, which is a type of arrythmia or abnormal heartbeat. Therefore, we typically want to quantify the amplitude and duration of the P-wave for AFib classification. The distance between the P-wave onset and onset of the QRS complex is the PR interval with a normal duration of 120-220 ms.

The QRS complex reflects depolarization of the left ventricle (since the electrical vector of the left ventricle is much larger than that of the right ventricle). A short QRS duration proves the ventricles are functioning properly and a broad QRS duration indicates that ventricular activation is slow and there could be a dysfunction in the electrical conduction system of the heart. The R-peak of the QRS complex is used to calculate the instantaneous heart rate from the interval between subsequent R-peaks (RR-interval). An RR-interval of 400 ms is equivalent to an instantaneous heart rate of 150 beats per minute ($60 s / 400 ms * 1000 ms / s$).

The ST segment is another important morphological feature of the ECG wave since ST elevation and depression are both associated with heart dysfunction like acute myocardial ischemia or ST-elevation myocardial infarction (STEMI). Elevation or depression are calculated as the difference (in millimeters) between the J point (where the ST segment starts) and the PR segment. Finally, the T-wave reflects a repolarization of the contractile cells and is also associated with a range o heart conditions.

Pan-Thompkins algorithm

The Pan-Thompkins algorithm is widely used and can be used for real-time continuous QRS detection. The algorithm is based on analysis of slope, width, and amplitude of ECG using a series of filters. An ECG signal first goes through a bandpass filter, then a differentiator, a squaring operation, a moving window integrator, and finally adaptive thresholding and search-back to find the R-peak.

Pan-Thompkins algorithm for QRS detection.

Raw ECG signals include muscle noise (from respiration), motion artifacts, the QRS complex, and P-T Waves. The band pass filter is designed to match the spectrum of QRS complex, attenuates muscle noise, 60Hz interference, baseline wander, and T-wave interference. Pass band of 5-15 Hz maximizes the QRS energy.

Filtering on waveforms can have negative effects. While low pass filters successfully reduce noise in ECG traces, they also reduce the QRS amplitude. High pass filters (e.g. low cutoff at 0.5 Hz) reduce baseline wander, but also introduce ST distortion. Using forward/backward filtering (highpass, reverse time, highpass, reverse time) removes most of the distortion introduced by high pass filters on the ST segment.

Here we use cascading filters combining a low pass filter and a high pass filter to mimic a bandpass filter. The filter attenuates the P and T waves (which peak at <5 Hz), which is a desired feature since the goal is to detect the QRS complex.

The filtered ECG is then differenced and squared to amplify the QRS complex. The derivative filter further suppresses low frequency components of P and T waves. Squaring makes the signal positive and enhances the derivatives by amplifying the high frequency QRS complex.

Next, the moving average filter over a 150ms window captures the duration of the QRS complex and gives us the integrated signal. This suppresses the smaller oscillations by smoothing out the residual high frequency components. Here we have to define an optimal window length for the moving average. Large windows merge the QRS and T waves together and small windows would produce several peaks at the QRS complex making it difficult to find the R-peak. In addition to detecting the QRS complex the moving average filter gives us the width of the QRS complex.

There are numerous heurestics for peak detection from the integrated signal (e.g. simple thresholding of the moving window integral). Pan-Thompkins proposes to use adpative thresholding and search-back to select a range of time values that correspond to QRS complexes by adapting to changes in ECG by computing running estimates of signal and noise peaks. Here instead I smooth the integrated signal with a gaussian filter to get the energy of the signal. The peak then corresponds to zero-crossings of the first difference where $x[i+1] < x[i]$.

import numpy as np
import scipy.signal

ecg = np.loadtxt("ecg.txt") # load the ECG signal
fs = 300 # Hz
tvec = np.arange(len(ecg)) / fs # time vector

max_QRS_duration = 0.150 # sec
low_cutoff = 5
high_cutoff = 15
window_size = int(max_QRS_duration * fs)

# apply a bandpass filter to the ECG signal
lowpass = scipy.signal.butter(1, high_cutoff / (fs / 2.0), "low")
highpass = scipy.signal.butter(1, low_cutoff / (fs / 2.0), "high")
ecg_low = scipy.signal.filtfilt(*lowpass, x=ecg)
ecg_band = scipy.signal.filtfilt(*highpass, x=ecg_low)

diff = np.diff(ecg_band)
squared = np.square(diff)

# moving average filter
# apply padding on both sides of the signal and convolve to get the integrated signal
mwa = np.pad(squared, (window_size - 1, 0), "constant", constant_values=(0, 0))
mwa = np.convolve(mwa, np.ones(window_size), "valid")
for i in range(1, window_size):
 mwa[i - 1] = mwa[i - 1] / i
mwa[window_size - 1 :] = mwa[window_size - 1 :] / window_size
mwa[: int(max_QRS_duration * fs * 2)] = 0

# smooth the moving window integrated signal with a gaussian filter and take the derivative
energy = scipy.ndimage.gaussian_filter1d(mwa, fs / 8.0)
energy_diff = np.diff(energy)

# peaks are the points where the derivative crosses zero, adjust window size
zero_crossings = (energy_diff[:-1] > 0) & (energy_diff[1:] < 0)
zero_crossings = np.flatnonzero(zero_crossings)
zero_crossings -= int(window_size / 2)

Stages of the beat detection algorithm. The vertical red line indicates the peaks of the QRS complex.

The figure above visualizes each step of the beat detection algorith. The vertical red line are all the zero-crossings (zero_crossings), which lines up with the peak of the ECG signal.

Other beat detection algorithms

Algorithm	Reference
Zeelenberg (1979)	Engelse, W.A.H., Zeelenberg, C (1979). A single scan algorithm for QRS detection and feature extraction, IEEE Comp. in Cardiology, vol. 6, pp. 37-42.
Pan (1985)	Pan, J., & Tompkins, W. J. (1985). A real-time QRS detection algorithm. IEEE transactions on biomedical engineering, (3), 230-236.
Hamilton (2002)	Hamilton, P. (2002, September). Open source ECG analysis. In Computers in cardiology (pp. 101-104). IEEE.
Zong (2003)	Zong, W., Heldt, T., Moody, G. B., & Mark, R. G. (2003). An open-source algorithm to detect onset of arterial blood pressure pulses. In Computers in Cardiology, 2003 (pp. 259-262). IEEE.
Christov (2004)	Ivaylo I. Christov, Real time electrocardiogram QRS detection using combined adaptive threshold, BioMedical Engineering OnLine 2004, vol. 3:28, 2004.
Elgendi (2010)	Elgendi, Mohamed & Jonkman, Mirjam & De Boer, Friso. (2010). Frequency Bands Effects on QRS Detection. The 3rd International Conference on Bio-inspired Systems and Signal Processing (BIOSIGNALS2010). 428-431.
Kalidas (2017)	Vignesh Kalidas and Lakshman Tamil (2017). Real-time QRS detector using Stationary Wavelet Transform for Automated ECG Analysis. In: 2017 IEEE 17th International Conference on Bioinformatics and Bioengineering (BIBE). Uses the Pan and Tompkins thresolding.
Nabian (2018)	Nabian, M., Yin, Y., Wormwood, J., Quigley, K. S., Barrett, L. F., Ostadabbas, S. (2018). An Open-Source Feature Extraction Tool for the Analysis of Peripheral Physiological Data. IEEE Journal of Translational Engineering in Health and Medicine, 6, 1-11.
Rodrigues (2021)	Rodrigues, Tiago & Samoutphonh, Sirisack & Plácido da Silva, Hugo & Fred, Ana. (2021). A Low-Complexity R-peak Detection Algorithm with Adaptive Thresholding for Wearable Devices.

Additional resources

Heart Rate Variability and Atrial Fibrillation

Wed, 13 Mar 2024 00:00:00 +0000

Atrial fibrillation (AFib) is a sustained cardiac arrythmia and is classified according to the temporal pattern of irregularly spaced heart beats. Patients with AFib have cardiac hemodynamic dysfunction, have up to 2-fold increase in risk of mortality, and a 6-fold increase in risk of stroke. The electrocardiographic presentation of AFib is continuous and rapid irregular electrical activity of the atria and absence of the P-wave because ventricular response is poorly coupled with atrial activity. These hallmark characteristic of AFib make ECG monitoring the most convenient tool to assist AFib diagnosis. Automated algorithms rely on one or more characteristics of the waveform including irregular rhythm, high-frequency chaotic atrial waveform, and absence of P waves. Measures of heart rate variability (HRV) and morphological analysis are the most common approaches.

R-R interval

We can use a beat detection algorithm to find the R-peaks of the QRS complexes, which are recorded as indices $R_{peaks}[n]$ for the $n$-th beat. We can convert the indices to an R-R interval (RRI) in units of seconds by taking the first difference and dividing by the sampling rate $f_s$:

$$RRI[n] = \frac{R_{peaks}[n] - R_{peaks}[n-1]}{f_s}$$

Poincare plots

A Poincare plot is a diagram in which the RR-interval ($RRI[n]$) is plotted as a function of the previous RR-interval ($RRI[n-1]$) and is a visual representation of heart rate variability. Patients with AFib have irregular RRI and the dispersion is wider. As such, we can quantify this dispersion by fitting an ellipsis and measuring the standard deviation along the major and minor axes. The area of the fitted ellipsis is also larger in AFib patients.

Javascript custom events

Mon, 08 Jan 2024 00:00:00 +0000

Custom events are a way to decouple a Javascript application. Trigger an event with the onclick handler (#!js onclick="triggerEvent(this,'my-event')") of an element, and listen for it elsewhere in the application.

function triggerEvent(el, name) {
 document.body.dispatchEvent(new CustomEvent('app:' + name, { detail: el }))
}
function main() {
 document.body.addEventListener('app:my-event', function (event) {
 console.log(event.detail) // element that triggered the event
 });
}
document.addEventListener('DOMContentLoaded', main)

Audio biosignal processing of phonocardiograms

Fri, 01 Jul 2022 00:00:00 +0000

A phonocardiogram (PCG) is a non-invasive assessment of the mechanical function of the heart. Cardiac auscultation and the analysis of the phonocardiogram can unveil fundamental clinical information regarding heart malfunctioning caused by congenital and acquired heart disease. This is achieved by detecting abnormal sound waves, or heart murmurs, in the PCG signal. This article uses data from the 2022 PhysioNet challenge to explore the spectral properties of PCG signals.

Below we see the time domain signal for a 5 sec window of PCG data. We can clearly see the S1 and S2 waves, which correspond to the beginning and end of the systolic phase of the heart beat, respectively.

The corresponding spectrogram shows the power at each frequency over time. The beats are clearly visible in both the time and frequency domains. There is some high frequency noise at about 1.5sec that we will next remove using a digital filter.

Spectrogram of the PCG signal

A Butterworth low-pass filter applied to this signal with 250 Hz cutoff frequency removes the high frequency noise at 1.5sec.

import scipy.signal

def butter(sig, freq, low_cutoff_hz, high_cutoff_hz, btype="low"):
 b, a = scipy.signal.butter(low_cutoff_hz, high_cutoff_hz / (freq / 2.0), btype)
 return scipy.signal.filtfilt(b, a, sig)

Most of the signal of interest is below 500Hz so low pass filtering removes the high frequency components and leaves the cardiac cycles intact. The figure below shows both the low pass filtered signal in the time domain (left) and frequency domain (right) on top and the high pass filtered signal on the bottom.

Low and high passed signals

Zooming in shows the cardiac cycles in more detail.

Original, low, and high passed signals

Time series similarity with random convolutional features and locality-sensitive hashing

Fri, 01 Jul 2022 00:00:00 +0000

Given a time series, like temperature readings from collection of sensors, we want to find sensors that have similar readings. This is a common problem in applications like sensor networks, IoT, and monitoring systems.

One approach is using random convolutional features to encode the signal and then use locality-sensitive hashing (LSH) to find similar signals. This approach is very fast and can be used to search through millions of time series signals. The time series neural hashing technique introduced here is a fast general-purpose search and retrieval algorithm.

Our specifications are:

We want to index a large number of time series signals and quickly retrieve similar time series in a way that scales computationally and has low storage requirements
Signals are of variable length with stretches of missing values. Imputation of missing values is not feasible
The representation should capture both the shape/structure and magnitude of the signal
Minimal feature engineering

Random convolutional neural hashing

Given a time series $X \in \mathbb{R}^{M \times N}$ of $M$ samples each of length $N$ containing missing values, we want to encode each sequence $x_m \in \mathbb{R}^{N}$ into a fixed length embedding vector that we can use for fast similarity search.

Normalize and clip $x_m$ in the range [0,1].
Fit a random convolutional encoder and save the parameters.
Embed $x_m$ to $e_m \in \mathbb{R}^{k}$ using the convolutional encoder.
Concatenate the sample stats $s_m={min, max, mean, P10, P25, P50, P90}$ from the normalized sequence to $e_m$ such that $[e_m; s_m] \in \mathbb{R}^{k+7}$
Weight sample stats $s_m$ by $\alpha$ and convolutional encodings $e_m$ by $1-\alpha$. ($\alpha=0.75$)
Define a seed matrix of shape $D \times k + 7$ from a standard normal distribution ($D=256$).
Calculate the hash string for each sample (e.g e9df77eb7c692f16)
Retrieval: Given a query hash, compute the hamming distance to find the nearest neighbor.

The random convolutional encoder is very fast to fit and evaluate. It’s generally good at capturing the structure of the signal however the magnitude may be lost because we use a max pooling layer that introduces scale and shift invariance. However, our similarity search should consider magnitude, which is important for our specific use case. Therefore, we introduce simple statistics about the magnitude and distribution of sample observations.

Pros:

Captures both signal structure (e.g. periodicity, stretches of missing data) and magnitude.
Fast and scalable both computationally and storage-wise.
Works on time series with missing values and of variable length.

Cons:

Sacrifices some precision since we are using random convolutional network where kernel weights, dilations, stide, are randomly selected.
Hashing-based retrieval is an approximate nearest neighbor approach which also has lower precision compared to exact nearest neighbor search or similarity search using tree-based methods with a distance metric.

Python project setup with Makefile and setup.py

Tue, 08 Mar 2022 00:00:00 +0000

All my Python projects are setup in the same way - as a Python package with a Makefile. This is a template for that setup.

$ tree .
.
├── Makefile
├── setup.py
├── myproject
│   ├── __init__.py

The setup.py file defines the requirements for the package, and the Makefile defines the commands to run.

The setup creates a virtual environment in .venv directory, installs the package in editable mode, and installs the development dependencies. Editable installs are useful when developing a package, as changes to the code are immediately available without having to reinstall the package. If you add a new dependency to the setup.py file, you will need to run make setup again to install it.

Here is an example setup.py file:

from setuptools import setup

VERSION = "0.1.0"

setup(
 # package name, which can be different from project, this is the name used
 # when installing the package with pip, e.g. pip install mypackage
 name="mypackage",
 version=VERSION,
 maintainer="",
 maintainer_email="",
 description="",
 license="",
 python_requires=">=3.10",
 # static files to include in the package
 package_data={
 "myproject": [
 "var/*",
 ]
 },
 # command line entry points
 entry_points={
 "console_scripts": [
 "myproject = myproject.__main__:main",
 ]
 },
 # fol
 packages=["myproject"],
 install_requires=[
 "pip",
 "pyarrow==13.0.0",
 "polars==0.18.15",
 "awscli==1.29.38",
 "boto3==1.28.38",
 "botocore==1.31.38",
 "python-dotenv==1.0.0",
 "uvicorn==0.24.0.post1",
 "requests==2.31.0",
 ],
 extras_require={
 "dev": [
 "ruff",
 "ipykernel",
 "pytest",
 ]
 },
 include_package_data=True,
 zip_safe=False,
 classifiers=[
 "Intended Audience :: Science/Research",
 "Programming Language :: Python :: 3.10",
 "Operating System :: OS Independent",
 ],
 keywords="python",
)

The Makefile looks like this:

help: ## Show this help
 @echo "\nSpecify a command. The choices are:\n"
 @grep -E '^[0-9a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf " \033[0;36m%-12s\033[m %s\n", $$1, $$2}'
 @echo ""
.PHONY: help

clean: ## Clean
 rm -rf ./.venv
 rm -rf ./dist
 rm -rf ./mypackage.egg-info
 rm -rf ./mypackage/__pycache__
 rm -rf ./myproject/*.so
 rm -rf ./myproject/__pycache__/
.PHONY: clean

setup: ## Editable install
 test -d .venv || python3 -m venv .venv
 . .venv/bin/activate; \
 python -m pip install --upgrade -i https://pypi.example.com/simple/ -e .[dev]
.PHONY: setup

server: ## Start local server
 uvicorn myproject.app:app --reload --workers=1 --reload-include="./myproject*"
.PHONY: server

Running make or make help will show the available commands:

Specify a command. The choices are:

 help Show this help
 setup Editable install
 server Start local server

Taxonomy of health data for machine learning

Wed, 23 Feb 2022 00:00:00 +0000

There is a wide variety of data types collected in the health system that can be utilized by machine learning models. These can include

patient-level information like demographics and socio-economic factors
hospital encounter-level information like admission source, ICU unit type, and discharge location
outcomes including diagnoses like billing codes and patient outcomes
interventions a patient received in the hospital like medications, invasive mechanical ventilation, oxygen support, pressors, fluids, blood transfusions, and ECMO
findings from radiological images, pathology images, and video recordings
laboratory measurements like blood gases, metabolic panels, liver panels, lipid panels, complete blood count, urinalysis, urine output, microbiology, and omics data
continuous waveforms like ECG, PPG, PCG, ABP, and etCO2 signals
nurse charted or automated vital sign collection including temperature, heart rate, blood pressure, and oxygen saturation
clinican and radiological notes

In-patient Data collected in the hospital is linked to patients using a unique medical record number (MRN). Data collected in out-patient settings, including at home, a nursing home, or in ambulatory care may not also be linked to the patients MRN. Even inside the hospital, linking waveforms (especially in time) with patient data in electronic health records is a significant challenge.

Cosine similarity 1D convolutions

Mon, 31 Jan 2022 00:00:00 +0000

The cosine similarity function below provides sign and scale independent 1D convolutions. It has a learnable parameter $p$ where large values of $p$ increase the sharpness of the cosine similarity. Here $u$ is the signal and $v$ is the kernel (e.g. $[1,2,3]$).

$$\text{CosineSim}(u,v)=\text{sign}(u \cdot v)\text{abs}(\frac{u \cdot v}{|u|_2 \cdot |v|_2})^{p^2}$$

We consider the kernel $[1,2,3]$ and a 1D time series signal with 5 distinct motifs:

A. exact match $[1,2,3]$
B. negative sign, exact match $[-1,-2,-3]$
C. downscaled exact match $[0.2, 0.4, 0.6]$
D. median $[2,2,2]$
E. reversed $[3,2,1]$

We can compare a 1D convolution with kernel_size=3, dilation = 1, and padding = (kernel_size-1)*dilation = 2 against the cosine similarity distance.

1D convolution correctly gives the largest activation to both of the exact matches (A and B). However, the convolution also gives a large activation to parts of the signal where there should not be a match. The median sequence of values B $[2,2,2]$ and the reversed sequence E $[3,2,1]$ get a significantly high activation despite having no similarity to the filter. The downscaled exact match C is not selected by the convolution because of the scale of the filter. Standard convolutions on raw (unnormalized) data are not scale and sign independent. Common normalizations, like batch or layer norm, calculates normalizing terms over samples or channels but not point-wise. The resulting convolved signal requires max pooling to find the subsequences of greatest correlation with the filter.

In contrast, the cosine similarity gives a score of 1 or -1 only for exact matches. The feature is detected independent of sign or scale. The figure below shows that the cosine similarity distance correctly detects the motifs A, B, and C. If we set the sharpness parameter to an arbitratily large value $p=9$ then the only points are the exact matches.

The output of CosineSim is clearly interpreted as the points of maximal correlation between the signal susequence and the filter, where filters represent subsequence templates.

def cosine_similarity(signal, kernel, sharpness, padding):
 """
 Compute the cosine similarity distance between a signal and a kernel.
 Outputs a sequence of the same length as the signal.

 Parameters
 ----------
 signal : torch.Tensor
 input [batch_size, channels, length]
 kernel : torch.Tensor
 filter [kernel_size]
 sharpness : float
 sharpness parameter
 padding : int
 padding size, (kernel size-1) * dilation

 Returns
 -------
 torch.Tensor
 output [batch_size, channels, length]
 """
 kernel_size = kernel.size(-1)
 x = F.pad(signal, (padding,padding))
 x = x.unfold(2, kernel_size, 1)
 sim = F.cosine_similarity(x, kernel, dim=-1)
 sgn = torch.sign(torch.einsum('bdij,k->bdi', x, kernel))
 sim = sgn * torch.pow(torch.abs(sim), (sharpness**2))
 sim = sim[:, :, : -padding].contiguous()
 return sim

Group-by and count in Numpy

Wed, 26 Jan 2022 00:00:00 +0000

The crosstab function takes a list of array-like objects and returns a contingency table of counts. A pure numpy implementation of a pivot-table like this is useful in environments where we don’t want to import the pandas package.

from typing import Tuple, List
import numpy as np

def crosstab(*args) -> Tuple[Tuple[np.ndarray], np.ndarray]:
 """
 Contingency table of counts.

 Parameters
 ----------
 args : list of array-like
 Arrays of discrete categorical data.

 Returns
 -------
 actual_levels : Tuple[np.ndarray]
 The actual levels of the categorical variables.
 count : np.ndarray
 The counts of the categorical variables cross-tabulated.

 Examples
 --------
 >>> categorical = [1,3,2,3]
 >>> covariate = [5,3,3,4]
 >>> levels, count = crosstab(categorical, covariate)
 """
 levels, indices = zip(*[np.unique(a, return_inverse=True) for a in args])
 count = np.zeros(list(map(len, levels)), dtype=int)
 np.add.at(count, indices, 1)
 return levels, count

Separable temporal convolutions

Sat, 22 Jan 2022 00:00:00 +0000

Given a multivariate time series $x \in \mathbb{R}^{B \times D \times T}$ with $D=3$ channels, $T=4$ timesteps and batch size $B=1$.

x = [
 [1,5,10,20],
 [100,150,200,250],
 [1000,1500,2000,2500],
]

# [batch_size=1, in_channels=3, timesteps=4]
xt = torch.FloatTensor(x)
xt = xt.unsqueeze(0)

The separable convolution learns a group of filters for each channel independently, without interactions across channels. Below, I define a 1D convolutional layer with kernel size 2 and learn 1 filter per channel (num_channels=1).

separable = True
in_channels = xt.shape[1]
num_channels = 1
kernel_size = 2
stride = 1
layer_i = 0
dilation_size = 2 ** layer_i
padding = (kernel_size - 1) * dilation_size
groups = in_channels if separable else 1
out_channels = in_channels * num_channels

For illustrative purposes, the weights are initialized to 1 and bias to 0.

conv1 = nn.Conv1d(
 in_channels,
 out_channels,
 kernel_size,
 stride=stride,
 padding=padding,
 dilation=dilation_size,
 groups=groups,
)
torch.nn.init.constant_(conv1.weight, 1)
torch.nn.init.constant_(conv1.bias, 0)

A separable convolution with kernel_size=2 and num_channels=1 is simply a weighted sum along each channel.

tensor([[[ 1, 6, 15, 30],
 [ 100, 250, 350, 450],
 [1000, 2500, 3500, 4500]]])

Increasing num_channels will learn a set of independent filters for each channel. For example, with num_channels=3 gives a total number of output channels of num_channels*in_channels=9.

tensor([[[ 1, 6, 15, 30],
 [ 1, 6, 15, 30],
 [ 1, 6, 15, 30],
 [ 100, 250, 350, 450],
 [ 100, 250, 350, 450],
 [ 100, 250, 350, 450],
 [1000, 2500, 3500, 4500],
 [1000, 2500, 3500, 4500],
 [1000, 2500, 3500, 4500]]])

When separable=False and num_channels=1, you get mixing between the channels:

tensor([[[1101, 2756, 3865, 4980],
 [1101, 2756, 3865, 4980],
 [1101, 2756, 3865, 4980]]])

Challenges in Machine Learning for Health

Wed, 01 Dec 2021 00:00:00 +0000

The secondary analysis of health data is challenging due to confounding, bias, uncertainty, and missingness.

Health data is irregularly sampled in time

Some data types are more frequently acquired than others. Vitals and laboratory measurmeents are taken for most in-patients. Clinical notes are also usually available for most patients. Ventilator settings, ECG signals, invasive blood pressure measurements, imaging data, and genomics data are are more infrequent. In the intensive care unit (ICU), for instance, continuous signals from ECG and ABP are typically sampled at 125Hz to 500Hz. The raw ECG, ABP, PPG signal can be downsampled to “high-frequency numerics” at 1sec or 1min interval. Some EHR databases will further apply a median filter so the highest sampling rate of vital signs like heart rate, blood pressure, temperature, and oxygen saturation is 5min. Ventilation settings and arterial blood gases can be charted every 6-12hours. Other metabolic and liver panels are often ordered every 12-24hours. These signals are acquired asynchronously, at irregular time intervals, and at different sampling rates.

Adherence to a standard nomenclature

EHR databases may not follow a standard nomenclature (like SNOMED, LOINC, MDIL, etc…), which introduces uncertainty in mapping measurements to standard concepts. OMOP is an example of an EHR database schema that attempts to heavily standardize concepts. However, the publicly accessible MIMIC and Philips eICU databases are not standardized and it is left up to the user to create concept mappings for medications, vital signs, laboratory measurements, and disease diagnoses. Many numeric fields like a laboratory measurement for creatinine can be entered as free text in the EHR user interface, which introduces further noise for the secondary analysis of EHR data.

Critical care data encodes physiology, clinical practice patterns, and clinicians concern

Electronic health records encode more than a patients physiology. The pattern of measurements, lab orders, and treatments capture the clinical decisions made at the bedside. The pattern of clinical decisions changes between institutions, between units in a hospital, the size of the institution (e.g. teaching vs community hospital). Disentangling physiology from clinical concern and care patterns is important to produce generalizable disease prediction models. However, this is not as important when using health data for operational research and to optimize hospital operations (e.g. forecasting bed occupancy).

Mapping patients across care settings

Connecting in-patient data collected in the hospital with out-patient data collected at home or in ambulatory care and emergency medical services is especially challenging. Even critical care data collected in the ICU can have periods of missing data where sensors like ECG, invasive ABP using an arterial line catheter, or PPG are missing for extended periods.

Data reliability

Nurse charted measurements can sometimes be unreliable. As an example, nurse charted respiratory rate at the bedside in the general ward is often rounded to a multiple of 5. Even continuously acquired signals can be unreliable. Timestamps between multiple sensors from different vendors (and sometimes the same vendor) can become desynchronized. The internal clock can drift. Some wearable devices may only save downsampled data instead of the raw signal.

Dataset and concept shift

Data drifts over time, institutions, hospital units, and countries. For example, ventilation settings and disease progression will change as more care givers adopt lung protective ventilation protocols; more patients have permissive hypertension and are on vasopressors in the neuro-ICU than in other units. Models that rely on interventions or physiological variables subject to dataset shift are susceptible to silent failure. Appropriate metrics to track dataset shift, model performance, and a pipeline to retrain and redeploy models are necessary.

Generalizability of machine learning models

Algorithms developed for the general ward (GW) setting may not translate to higher acuity settings, and vice versa. Most patients in the GW are not on continuous monitoring and instead vitals are aperiodically nurse charted. Differences in the severity of illness and treatment patterns are also significantly different. Models should be carefully deployed in settings where they were designed and tested.

Reproducibility

Reproducibility of the data extraction pipeline is important for documentation and verification. This becomes challenging for a large team working on different parts of the data extraction pipeline with tasks like concept mapping, data labeling, cohort selection, and data pre-processing divided among scientists. Start by coordinating on data storage, naming conventions (for code and datasets), and tooling. I have found that documenting the pipeline in single function call helps with posterity and reproducibility.

Actionability

Algorithms that are actionable are tied to a protocol or clinical decision support system that is integrated into the clinical workflow. For example, predicting the that a patient is at high risk of hemodynamic shock on it’s own is not useful unless it is tied to a protocol that triggers an alert and a clinical decision support system that recommends appropriate fluids or pressors.

Quantile binning with missing data

Sun, 03 Oct 2021 00:00:00 +0000

This uses Numpy and Numba for fast binning of numerical data to quantiles. It also supports missing data.

import numpy as np
from numba import njit

def _find_binning_thresholds(data, max_bins=256, subsample=int(2e5)):
 if not (2 <= max_bins <= 256):
 raise ValueError(
 "max_bins={} should be no smaller than 2 "
 "and no larger than 256.".format(max_bins)
 )

 percentiles = np.linspace(0, 100, num=max_bins + 1)
 percentiles = percentiles[1:-1]
 binning_thresholds = []
 for f_idx in range(data.shape[1]):
 col_data = np.ascontiguousarray(data[:, f_idx], dtype=np.float64)
 mask = np.isfinite(col_data)
 col_data = col_data[mask]
 distinct_values = np.unique(col_data)
 if len(distinct_values) <= max_bins:
 midpoints = distinct_values[:-1] + distinct_values[1:]
 midpoints *= 0.5
 else:
 midpoints = np.percentile(
 col_data, percentiles, interpolation="midpoint"
 ).astype(np.float64)
 binning_thresholds.append(np.unique(midpoints))
 return binning_thresholds


@njit()
def _map_num_col_to_bins(data, binning_thresholds, binned):
 for i in range(data.shape[0]):
 left, right = 0, binning_thresholds.shape[0]
 while left < right:
 middle = (right + left - 1) // 2
 if data[i] <= binning_thresholds[middle]:
 right = middle
 else:
 left = middle + 1
 binned[i] = left


def _map_to_bins(data, binning_thresholds, binned):
 """Bin numerical values to discrete integer-coded levels."""
 for feature_idx in range(data.shape[1]):
 _map_num_col_to_bins(
 data[:, feature_idx],
 binning_thresholds[feature_idx],
 binned[:, feature_idx],
 )


def _assign_nan_to_bin(binned, X, actual_n_bins, assign_nan_to_unique_bin=False):
 mask = np.isnan(X)
 for i in range(X.shape[1]):
 binned[mask[:, i], i] = actual_n_bins[i] if assign_nan_to_unique_bin else np.nan
 return binned


class QuantileBinning():
 def __init__(self):
 self.bin_thresholds = []
 self.n_bins = []

 def fit(self, X):
 self.bin_thresholds = _find_binning_thresholds(X)
 self.n_bins = np.array(
 [thresholds.shape[0] + 1 for thresholds in self.bin_thresholds], dtype=np.uint32
 )

 def transform(self, X, assign_nan_to_unique_bin=False):
 binned = np.zeros_like(X, dtype=np.float32, order="F")
 _map_to_bins(X, self.bin_thresholds, binned)
 binned = _assign_nan_to_bin(binned, X, self.n_bins, assign_nan_to_unique_bin)
 return binned

Ensemble decision trees in Numba

Tue, 21 Sep 2021 00:00:00 +0000

Representing ensembles of decision trees using numpy arrays with fast numba operations.

import math
import numpy as np
from numba import njit, prange


@njit
def take(X, inds):
 """Multidimensional indexing for numba"""
 n = len(X)
 y = np.zeros(n, dtype=X.dtype)
 for i in range(n):
 y[i] = X[i, inds[i]]
 return y


@njit
def next_node(node_id, value, thr):
 """A vectorized operation to find the next node in a binary tree given a
 value and threshold.
 Args:
 node_id: int or array of the current node, root node_id = 0
 Return:
 new node id, left node when value <= thr, right node when value > thr
 """
 return (node_id << 1) + 1 + (1 * (thr < value))


@njit
def leaf(X, features, thresholds, reset_leaf_index=1):
 """Find the leaf node index along a decision path given a tree feature
 indices, thresholds, and design matrix.
 Args:
 X: 2D design matrix of shape [nsamples, nfeatures]
 features: feature indices as dtype np.int32 of shape [internal_nodes]
 thresholds: split thresholds as dtype np.float64 of shape [internal_nodes]
 reset_leaf_index: returns leaf index initialized from 1
 """
 nsamples = len(X)
 node_id = np.zeros(nsamples, dtype=np.int64)
 n = len(features)
 depth = int(math.log(n+1)/math.log(2))
 internal_nodes = 2**(depth) - 1
 for i in range(depth):
 feature_ind = features[node_id]
 value = take(X, feature_ind)
 thr = thresholds[node_id]
 node_id = next_node(node_id, value, thr)
 if reset_leaf_index == 1:
 node_id = node_id - internal_nodes
 return node_id


@njit
def leaf_tokens(X, trees, nleaves_per_tree):
 """Tokenize a design matrix with leaf indices.
 Args:
 X: 2D design matrix of shape [nsamples, nfeatures]
 trees: 3D matrix deifining an ensemble of decision trees of shape
 [ntrees, internal_nodes, 2]
 nleaves_per_tree: number of leaves in each decision tree, 2**depth
 """
 nsamples = X.shape[0]
 ntrees = trees.shape[0]
 leaves = np.zeros((nsamples,ntrees), dtype=np.int64)
 for i in prange(0,ntrees):
 features = trees[i,:,0].astype(np.int64)
 thresholds = trees[i,:,1]
 leaves[:,i] = leaf(X, features, thresholds) + nleaves_per_tree * i
 return leaves


@njit
def random_ensemble_decision_trees(ntrees, depth, nfeatures):
 """Generate a random ensemble of decision trees
 Args:
 ntrees: number of trees in ensemble
 depth: height of each tree
 nfeatures: number of features in design matrix
 Returns:
 trees: 3D matrix of shape [ntrees, internal_nodes, 2] where internal_nodes
 is the number of non-leaf nodes (including root node) calculated
 as 2^{depth} - 1 and the last dimension includes feature index and
 splitting thresholds
 """
 internal_nodes = 2**(depth) - 1
 total_internal_nodes = ntrees * internal_nodes
 trees = np.zeros((total_internal_nodes, 2)) # [features, thresholds]
 trees[:,0] = np.random.randint(0, nfeatures, total_internal_nodes)
 trees[:,1] = np.random.rand(total_internal_nodes)
 trees = trees.reshape(ntrees, internal_nodes, -1) # [ntrees, internal_nodes, 2]
 return trees


nsamples, nfeatures = 1000, 30
X = np.random.rand(nsamples,nfeatures)

ntrees = 1000
depth = 1
nleaves = 2**depth
total_nleaves = ntrees * nleaves
trees = random_ensemble_decision_trees(ntrees=ntrees, depth=depth, nfeatures=nfeatures)
leaves = leaf_tokens(X, trees, nleaves)

Feature engineering for time series data using Numba

Tue, 07 Sep 2021 00:00:00 +0000

Feature engineering over a multivariate time series with missing data.

Given a sequence of measurements values: np.ndarray and observation times times: np.ndarray, we want to engineer features for a machine learning model that captures the temporal trends and statistics over different temporal windows. Our data is irregularly sampled and have missing data.

Features representing the magnitude, dispersion, direction of change, and temporal trends are derived. Every time point is represented with 4 categories of features:

Magnitude of the most recent observation within the last 6h for vital signs and 24h for laboratory measurements.
Dispersion measured as the range over a short and long-time window.
Direction of change (increasing, decreasing, no change) over a short and long-time window.
Exponential moving averages (EMA) with varying decay rates that specify how the much impact each past observation has on the current mean. EMA features were calculated on the forward filled magnitudes and using an EMA algorithm specifically for irregularly sampled time series.

The engineered features include:

dt: time elapsed since the measurement was made
val: most recent measurement, measurements are forward filled up to a maximum duration after
srng: range as $\frac{x_{max}-x_{min}}{x_{max}+x_{min}}100$ over a short window
ssgn: sign of the change (-1 or +1) between the first and last measurement in a short window, windows without measurements are filled with 0
lrng: range as (max-min)/(max+min) * 100 over a long window
lsgn: sign of the change (-1 or +1) between the first and last measurement in a long window, windows without measurements are filled with 0
sema: slow exponential moving average, calculated after forward filling
fema: fast exponential moving average, calculated after forward filling

We also need to define the following settings for every variable:

forward_fill_duration: duration in minutes to forward fill a missing variable
short_rolling_window_size: short rolling window size in minutes
long_rolling_window_size: long rolling window size in minutes
slow_ema_tau: decay rate for slow exponential moving average in minutes
fast_ema_tau: decay rate for fast exponential moving average in minutes

Some derived features have no missing data, including dt, ssgn, srng, lsgn, and lrng. Other variables like val, sema, fema have missing values encoded as NaN. Missing data is forward filled up to forward_fill_duration and all missing values after the forward filling time are NaN.

from numba import njit, prange

@njit()
def ema(times, values, tau):
 """Exponential moving average for irregularly sampled time series.
 Units for ``times`` and ``tau`` should be in hours.

 Args:
 times (np.ndarray): 1D array of measurement times in hours
 values (np.ndarray): 1D array of measurements
 tau (float): time decay in hours
 """
 n = len(values)
 ret = np.empty(n, dtype=np.float64) * np.nan
 ret[0] = values[0]
 last_i = 0
 for i in range(1, n):
 if np.isnan(times[i]) | np.isnan(values[i]):
 continue
 alpha = (times[i] - times[last_i]) / tau
 w = np.exp(-alpha)
 if alpha > 1e-6:
 w2 = (1 - w) / alpha
 else:
 # use Taylor expansion for numerical stability
 w2 = 1 - (alpha / 2) + (alpha * alpha / 6) - (alpha * alpha * alpha / 24)
 ret[i] = (ret[last_i] * w) + (values[i] * (1 - w2)) + (values[last_i] * (w2 - w))
 last_i = i
 return ret


@njit()
def ffill(times, values):
 """Forward fill an array of values and times.
 Times indicate the time of last observation.

 Args:
 arr (np.ndarray): array of values with nan
 times (np.ndarray): array of times

 Returns:
 2D array with forward filled times (index 0)
 and values (index 1).
 """
 out_values = values.copy()
 out_times = times.copy()
 n = values.shape[0]
 out = np.zeros((n, 2), dtype=np.float64)
 lastval = np.nan
 lasttime = np.nan
 for row_idx in range(n):
 if np.isfinite(values[row_idx]):
 lastval = values[row_idx]
 lasttime = times[row_idx]
 out_values[row_idx] = lastval
 out_times[row_idx] = lasttime
 out[:, 0] = out_times
 out[:, 1] = out_values
 return out


@njit()
def sliding_windows(times, window_size):
 tdiff = times - times.reshape(-1, 1)
 twindows = (tdiff >= 0) & (tdiff <= window_size)
 return twindows


@njit()
def feature_engineering(
 times,
 values,
 forward_fill_duration,
 short_rolling_window_size,
 long_rolling_window_size,
 slow_ema_tau,
 fast_ema_tau,
):
 """Feature engineering.

 1. Age of the observation
 2. Most recent forward-filled measurement
 3. Dispersion as (max-min)/(max+min) over short window
 4. Sign of change between first and last value in short window
 5. Dispersion as (max-min)/(max+min) over long window
 6. Sign of change between first and last value in long window
 7. slow EMA
 8. fast EMA

 Forward filling replaces nan values up to a maximum duration given by
 `forward_fill_duration`. Any remaining missing values should be imputed
 with the median or by sampling from a standard reference range.

 Args:
 values (np.ndarry): 2D array of observations
 times (np.ndarry): 1D array of observation times
 forward_fill_duration (float): duration to forward fill missing values
 short_rolling_window_size (float): observation duration for rolling
 statistics with a short lookback window
 long_rolling_window_size (float): observation duration for rolling
 statistics with a long lookback window
 slow_ema_tau (float): slow decay rate weights more of the past
 fast_ema_tau (float): fast decay rate weights more of the present

 Returns:
 2D matrix with shape [n_samples, 8], where the features are:
 suffixes = ['dt', 'val', 'short_rng', 'short_sgn', 'long_rng', 'long_sgn', 'slow_ema', 'fast_ema']
 """
 DT_IX = 0
 VAL_IX = 1
 SRNG_IX = 2
 SSGN_IX = 3
 LRNG_IX = 4
 LSGN_IX = 5
 SEMA_IX = 6
 FEMA_IX = 7

 n = values.shape[0]
 derived = np.zeros((n, 8)) * np.nan

 missing = np.isnan(values)

 # return a median imputed array if there are no observations
 if np.all(missing):
 derived[:, SSGN_IX] = 0.
 derived[:, LSGN_IX] = 0.
 derived[:, SRNG_IX] = 0.
 derived[:, LRNG_IX] = 0.
 return derived

 # indicate samples within a fixed window for each sample
 short_windows = sliding_windows(times, short_rolling_window_size)
 short_windows = short_windows & ~missing.reshape(-1,1)
 long_windows = sliding_windows(times, long_rolling_window_size)
 long_windows = long_windows & ~missing.reshape(-1,1)

 # forward fill
 xfill = ffill(times, values)
 dt = times - xfill[:,0]

 # 0. Age of the measurement
 derived[:, DT_IX] = dt

 # 1. last measured value
 derived[:, VAL_IX] = xfill[:, 1]
 expired = ~np.isnan(dt)
 expired[expired] = dt[expired] >= forward_fill_duration
 derived[expired, VAL_IX] = np.nan

 # rolling statistics
 for t in range(n):
 # samples in this window
 short_inds = short_windows[:, t]
 long_inds = long_windows[:, t]
 short_vals = values[short_inds]
 long_vals = values[long_inds]

 if short_vals.size > 0:
 # 2. rolling dispersion
 high = np.max(short_vals)
 low = np.min(short_vals)
 total = high + low
 if total > 1e-6:
 derived[t, SRNG_IX] = (high - low) / (total + 1e-6) * 100
 else:
 derived[t, SRNG_IX] = 0.

 # 3. sign of change between last and first value
 if short_vals.size >= 2:
 derived[t, SSGN_IX] = np.sign(short_vals[-1] - short_vals[0])
 else:
 derived[t, SSGN_IX] = 0.
 else:
 derived[t,SRNG_IX] = 0.
 derived[t,SSGN_IX] = 0.

 if long_vals.size > 0:
 # 4. rolling dispersion
 high = np.max(long_vals)
 low = np.min(long_vals)
 total = high + low
 if total > 1e-6:
 derived[t, LRNG_IX] = (high - low) / (total + 1e-6) * 100
 else:
 derived[t, LRNG_IX] = 0.

 # 5. sign of change between last and first value
 if long_vals.size >= 2:
 derived[t, LSGN_IX] = np.sign(long_vals[-1] - long_vals[0])
 else:
 derived[t, LSGN_IX] = 0.
 else:
 derived[t, LRNG_IX] = 0.
 derived[t, LSGN_IX] = 0.

 # exponential moving average using the forward filled data
 ema_slow = np.empty(n, dtype=np.float64) * np.nan
 ema_fast = np.empty(n, dtype=np.float64) * np.nan
 m = np.isfinite(derived[:, 1])
 vals = derived[m, 1]
 if vals.size > 0:
 ema_slow = ema(times[m], derived[m, 1], slow_ema_tau)
 ema_fast = ema(times[m], derived[m, 1], fast_ema_tau)
 # 6. slow EMA
 derived[m, SEMA_IX] = ema_slow
 # 7. fast EMA
 derived[m, FEMA_IX] = ema_fast

 return derived

Select a random window of maximum duration in NumPy

Mon, 02 Aug 2021 00:00:00 +0000

Given a sequence of indices $x$ and times $t$, we want to select a window of data from $x$ with a maximum duration and a maximum number of measurements.

def select_random_window(
 x: np.ndarray, t: np.ndarray, max_window_size: int, dur: float
):
 """Select a random window of consecutive elements with a maximum duration.
 This function has extra logic that ensures the selected windows have length
 max_window_size whenever possible.

 Args:
 x (np.ndarray): sequence of indices
 max_window_size (int): maximum number of observations

 Returns:
 1D array of samples in the observation window.
 """
 L = len(x)
 if L <= max_window_size:
 return x
 if (t[-1] - t[0]) <= dur:
 return x
 L_tmp = L
 if L < max_window_size - 6:
 L_tmp = L - int(max_window_size / 2)
 else:
 L_tmp = L - 1
 start = random.randint(0, L_tmp)
 t_start = t[start]
 t_end = t_start + dur
 if t_end > t[-1]:
 return x[start:]
 end = np.argwhere((t - t_end) >= 0)[0][0]
 start = start if end <= L else max(0, start - (end - L))
 return x[start:end]

Causal inference learners

Mon, 01 Feb 2021 00:00:00 +0000

Given covariates $X$, treatment indicator $W$ and a binary outcome $Y \in {0,1}$.

Inverse probability of treatment weights

Fit a model on observations $X$ to predict the treatment $W$, $p_{w_i}(x_i)=P(W=w_i|X=x_i)$ and use the probability of being treated as a sample weight to predict the outcome $Y$.

$$IPTW_i=\frac{1}{p_{w_i}}$$

S-learner

A single model is trained on observations and treatments.

$$\begin{aligned} \hat{\mu} = M(Y \sim (X,W))\\ \hat{\tau}(x)=\hat{\mu}(x,1) - \hat{\mu}(x,0) \end{aligned}$$

T-learner

Two models, one model per treatment, are trained on the observations only.

$$\begin{aligned} \hat{\mu}_0 = M_0(Y^0 \sim X^0)\\ \hat{\mu}_1 = M_1(Y^1 \sim X^1)\\ \hat{\tau}(x)=\hat{\mu}_1(x)-\hat{\mu}_0(x) \end{aligned}$$

X-learner

The X-learner first estimates the response function ($\hat{\mu}0$ and $\hat{\mu}1$), then coomputes the imputed treatment effects ($\tilde{D}{i}^{1}$ and $\tilde{D}{i}^{0}$), then estimates the conditional average treatment effect for treated and controls ($\hat{\tau}_1$ and $\hat{\tau}_0$), and finally averages the estimates.

$$\begin{aligned} \hat{\mu}0 = M_1(Y^0 \sim X^0)\\ \hat{\mu}1 = M_2(Y^1 \sim X^1)\\ \tilde{D}{i}^{1}=Y_i^1 - \hat{\mu}{0}(X_i^1)\\ \tilde{D}_{i}^{0}=\hat{\mu}1(X_i^0) - Y{i}^{0}\\ \hat{\tau}_1=M_3(\tilde{D}^{1} \sim X^{1})\\ \hat{\tau}_0=M_4(\tilde{D}^{0} \sim X^{0})\\ \hat{\tau}(x)=g(x)\hat{\tau}_0(x)+(1-g(x))\hat{\tau}_1(x) \end{aligned}$$

$g(x)\in[0,1]$ is a weighting function which is chosen to minimize the variance of $\hat{\tau}(x)$. $g(x)$ can be estimated by the propsensity score or set to a constant and equal o the ratio of treated to untreated samples.

G-Computation

Parameters we can derive from the G-Computation:

For a binary outcome, $Y\in{0,1}$, $P(Y=1)=E[Y]$.
$E[Y^1]-E[Y^0]$ is the causal risk difference due to treatment
$E[Y^1]/E[Y^0]$ is the causal relative risk
$EY^1/EY^0$ is the causal odds ratio due to treatment

Cloudflare pages for static hosting

Thu, 21 Jan 2021 00:00:00 +0000

The asifr.com website has been hosted on a cheap DigitalOcean server for the last 8 years. I hadn’t touched the server for many years and it stagnated running Ubuntu 14 and didn’t have HTTPS. Today, I migrated the static site to Cloudflare Pages for free hosting in two steps:

Create a github repository with the static files. The repo can be public or private.
Connect Cloudflare pages to the github repo and set the path to the static files in your repo (e.g. /public or /dist).

Offline, I use pandoc to compile a folder of Markdown files into static HTML files.

Cloudflare Pages automatically rebuilds the site when I commit the changed files to the github repo. However, since we don’t have a cloud-based build step and are generating the static files offline, Cloudflare Pages simply serves the public static files.

Cloudflare gives a nice domain name: asifr.pages.dev, but also provides custom DNS service so I can point asifr.com to asifr.pages.dev. It takes a few hours for the domain name propagation once the nameservers are transferred and the CNAME is created. Now asifr.com has secure HTTPS connections and is served from Cloudflare’s edge network.

The benefits are: I get a free hosting service since Cloudflare Pages provides unlimited bandwidth, requests, and sites (the limitation of the free-tier is I get 500 builds per month – which is plenty for a simple personal site that is updated infrequently). You also get security, HTTPS, analytics, and users get fast access to the site from anywhere in the world.

Interpretable univariate risk curves

Sun, 05 Jul 2020 00:00:00 +0000

High accuracy complex models, like neural networks and generalized additive models, come at the expense of interpretability. The contribution of individual features to the outcome are difficult to understand in complex models and have recieved significant criticism, partcularly in high-stakes settings (like medicine and criminal cases) where trust is critical. Intelligible models develop trust with the user and helps us debug counter-intuitive or inaccurate relationships learned by the model. Before applying complex models, it’s usually a good idea to look at odds ratios and risk curves to visualzie how each variable is related to the outcome. Odds ratios can quantify the strength of the relationship and is a good first step in ranking features that have a strong association with the outcome, compared to features with weaker associations. Risk curves help us visualize the shape of the association (e.g. linear, non-linear, monotonic, increasing, decreasing).

Odds-ratios

Odds ratios are widely used to compare the relative odds of the occurrence of the outcome (e.g. disease), given exposure to a feature. The odds ratio can also be used to determine whether a particular feature is a risk factor for a an outcome, and to compare the magnitude of various risk factors for that outcome.

OR = 1 Feature does not affect odds of outcome
OR > 1 Feature is associated with higher odds of outcome
OR < 1 Feature is associated with lower odds of outcome

import statsmodels.api as sm
from scipy import stats
stats.chisqprob = lambda chisq, df: stats.chi2.sf(chisq, df)

def univariatelr(X,Y,feature_name,binary=False):
 res = sm.Logit(Y, X, missing='drop').fit(disp=0, intercept=True)
 ci = np.round(np.exp(res.conf_int(alpha=0.05, cols=None)),2).squeeze()
 r = {'feature':feature_name,
 'nobs': np.sum(X==1) if binary else res.nobs,
 'coeff': round(res.params[0],2),
 'OR': round(np.exp(res.params[0]),2),
 'ci_low': ci[0],
 'ci_high': ci[1],
 'p-value': np.round(res.pvalues[0],2)}
 return r

The coefficients of a logistic regression model are the log-odds and taking the exponential of the coefficient gives us the odds-ratio. Odds-ratios are plotted along with the confidence interval. If the confidence interval overlaps with OR=1, than the feature has a weak or no association with the outcome. Below is an example of the odds-ratios from univariate logistic regressions where the outcome label is survived or expired in the ICU.

Risk curves

Risk curves describe the relationship between a feature and an outcome. For example a risk curve can be used to understand the relationship between a physiological measure like heart rate and an outcome like the probability of developing a disease. These graphical interpretations help explain hidden relationships in data, are interpretable by a non-technical audience, and relatively easy to construct.

The risk curve for lactate shows that the risk of mortality monotonically increases with increasing lactate. We could also reasonably draw a cutoff around Lactate > 2 mmol to and group these high lactate patients for further analysis.

Building such a risk curve is simple enough. Given a table of feature value and label:

Value	Label
0.800	0
4.767	1
0.700	0
1.100	0
1.575	0

First bucket each sample into a bin, for example: (0.1 , 0.6), (0.6, 1.2), …, (4.4, 5.). You can use quantiles to determine these bins or just equally divide your data into bins. We can use the histogram function to divide our data into equal sized bins with a specified range and count the number of samples that fall into each bin (h_control, h_treated).

x = df.Value
y = df.Label
x = x[(x>=minval) & (x<=maxval)]
xmin = x.min()
xmax = x.max()
h_control, b_control = np.histogram(x[y==0], range=(xmin, xmax), bins=10)
h_treated, b_treated = np.histogram(x[y==1], range=(xmin, xmax), bins=10)
risk = np.log((h_treated/h_treated.sum())/(h_control/h_control.sum()))
vals = b_unstable[:-1]
plt.plot(vals, risk)

The log-odds (risk) is simply the fraction of treated samples in the bin divided by the fraction of untreated samples in the bin.

Transform Grouped Pandas DataFrame to Numpy Array

Mon, 22 Jun 2020 00:00:00 +0000

This snippet transforms a tall Pandas DataFrame with time-series data into a Numpy array while preserving the grouping. This is a common use case for me when preparing training data for recurrent neural networks, where each training sample belongs to a group (EventID below), feature values (FeatureValue) are orded by time (DateTime), and I want to get the length of each sample (needed to train an RNN with variable length sequences).

EventID	DateTime	FeatureValue
1	0	80
1	5	90
2	0	75
2	10	80

event_col = 'EventID'
time_col = 'DateTime'
value_col = 'FeatureValue'
xt = df.loc[:,[time_col, value_col]].values
g = df.reset_index(drop=True).groupby(event_col)
xtg = [xt[i.values,:] for k,i in g.groups.items()]
SignalLengths = [len(i.values) for k,i in g.groups.items()]
X_signal = np.array(xtg)
EventIDs = list(g.groups.keys())

AWS Lambda Web Scraper

Sun, 21 Jun 2020 00:00:00 +0000

This is a small AWS Lambda function to scrape websites using axios and store the data in a MongoDB document. You can setup an API Gateway to the Lambda function and use GET requests to call the function.

Features

Randomly selects from a set of headers with each call.
Automatically sets the host and referer to the same domain.
Saves the response to MongoDB.
Optionally sets the header to json if you expect the output to be json format
Optionally sets the request to XMLHttpRequest

Install the required Node modules: npm install axios mongodb dotenv

const request = require('axios');
const MongoClient = require('mongodb').MongoClient;
const crypto = require('crypto');

// local .env files are loaded into process.env
require('dotenv').config({silent: false});

// load the MongoDB connection string from the .env file
const mongo_host = process.env.MONGO

// database and collection name
databaseName = 'scraper'
collectionName = 'rawdata'

// set of headers from which we will randomly select
let headers_list = [
 // Firefox 77 Mac
 {
 "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:77.0) Gecko/20100101 Firefox/77.0",
 "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
 "Accept-Language": "en-US,en;q=0.5",
 "Referer": "https://www.google.com/",
 "DNT": "1",
 "Connection": "keep-alive",
 "Upgrade-Insecure-Requests": "1"
 },
 // Firefox 77 Windows
 {
 "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:77.0) Gecko/20100101 Firefox/77.0",
 "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
 "Accept-Language": "en-US,en;q=0.5",
 "Accept-Encoding": "gzip, deflate, br",
 "Referer": "https://www.google.com/",
 "DNT": "1",
 "Connection": "keep-alive",
 "Upgrade-Insecure-Requests": "1"
 },
 // Chrome 83 Mac
 {
 "Connection": "keep-alive",
 "DNT": "1",
 "Upgrade-Insecure-Requests": "1",
 "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/83.0.4103.97 Safari/537.36",
 "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9",
 "Sec-Fetch-Site": "none",
 "Sec-Fetch-Mode": "navigate",
 "Sec-Fetch-Dest": "document",
 "Referer": "https://www.google.com/",
 "Accept-Encoding": "gzip, deflate, br",
 "Accept-Language": "en-GB,en-US;q=0.9,en;q=0.8"
 },
 // Chrome 83 Windows
 {
 "Connection": "keep-alive",
 "Upgrade-Insecure-Requests": "1",
 "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/83.0.4103.97 Safari/537.36",
 "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9",
 "Sec-Fetch-Site": "same-origin",
 "Sec-Fetch-Mode": "navigate",
 "Sec-Fetch-User": "?1",
 "Sec-Fetch-Dest": "document",
 "Referer": "https://www.google.com/",
 "Accept-Encoding": "gzip, deflate, br",
 "Accept-Language": "en-US,en;q=0.9"
 }
]

function isValidURL(string) {
 var res = string.match(/(http(s)?:\/\/.)?(www\.)?[-a-zA-Z0-9@:%._\+~#=]{2,256}\.[a-z]{2,6}\b([-a-zA-Z0-9@:%_\+.~#?&//=]*)/g);
 return (res !== null)
};

module.exports.scrape = async event => {
 // start by parsing the body assuming a POST statement with a JSON body
 let body = JSON.parse(event.body)

 // url is required
 if (!('url' in body)) {
 return {ok: 0, msg: 'Missing URL'}
 }

 // check the url is valid
 if (!isValidURL(body.url)) {
 return {ok: 0, msg: 'Invalid URL'}
 }

 let url = body.url
 let host = new URL(url)

 // randomly select a header
 let headers = headers_list[Math.floor(Math.random() * headers_list.length)]
 // the request should look like it is originating from the host
 headers['Host'] = host.host
 // referer is from the same domain, referers from google.com are often
 // redirected, which we want to avoid
 headers['Referer'] = host.origin

 // set json headers if we expect the response to be in json
 if ('json' in body) {
 headers['Accept'] = 'application/json, text/javascript, */*; q=0.01'
 }

 // set XMLHttpRequest header, which helps when calling private APIs that
 // would typically be loaded by AJAX calls
 if ('ajax' in body) {
 headers['X-Requested-With'] = 'XMLHttpRequest'
 }

 // send a GET request with our headers
 const response = await request({
 'url': url,
 'method': 'get',
 'headers': headers,
 });

 if (response.status == 200) {
 // create a data object containing the response body and headers
 let date = new Date()
 let data = {
 'url': url,
 'url_hash': crypto.createHash('md5').update(url).digest("hex"),
 'host': host.host,
 'data': response.data,
 'processed': false,
 'scraped_at': date,
 'scraped_year': date.getFullYear(),
 'scraped_month': date.getMonth() + 1,
 'scraped_day': date.getDate(),
 'response_headers': response.headers,
 'request_headers': headers
 }

 // create a connection to the MongoDB
 const client = await MongoClient.connect(mongo_host, {useUnifiedTopology: true});
 // select the database
 const db = client.db(databaseName);
 // insert data into collection in database
 let r = await db.collection(collectionName).insertOne(data);
 // close the connection to MongoDB
 client.close();

 if (r.insertedCount == 1) {
 // return the newly created ObjectID if a new document was successfully inserted
 return {
 ok: 1,
 url: url,
 insertedId: r.insertedId,
 };
 }
 } else {
 return {
 ok: 0,
 url: url,
 status: response.status,
 msg: 'Bad response status'
 };
 }

 return {
 ok: 0,
 url: url
 };
};

Deploy static website with rsync

Sun, 21 Jun 2020 00:00:00 +0000

Create a Makefile and use rsync to deploy your ./public folder to a remote server. Replace <IP-ADDRESS> with the website location and set FOLDER to the remote location where you want to upload your local files. In practice it’s nice to have a staging folder to test out your website before deployment. Use make staging and make deploy to sync changes.

.PHONY: staging deploy

SERVER = <IP-ADDRESS>
FOLDER = /var/www/html
STAGING = /var/www/html/staging
USER = root

deploy:
 rsync -zarvh ./public/* --compress --recursive --checksum --delete --itemize-changes --exclude-from exclude.rsync $(USER)@$(SERVER):$(FOLDER)

staging:
 rsync -zarvh ./public/* --compress --recursive --checksum --delete --itemize-changes --exclude-from exclude.rsync $(USER)@$(SERVER):$(STAGING)

Svelte Webpack Boilerplate

Fri, 22 May 2020 00:00:00 +0000

This setup uses webpack to bundle code and Svelte to create the user interface. This looks for files under /src and saves the compiled and minified javascript code in /public/js.

Install the required Node modules: npm install svelte svelte-loader.

const path = require('path');

module.exports = {
 entry: {
 dashboard: './src/dashboard.js',
 account: './src/account.js',
 },
 output: {
 path: path.resolve(__dirname, 'public/js/'),
 filename: "[name].js"
 },
 mode: "production",
 module: {
 rules: [
 {
 test: /\.(html|svelte)$/,
 exclude: /node_modules/,
 use: {
 loader: 'svelte-loader',
 options: {}
 },
 }
 ]
 },
 plugins: [],
 resolve: {
 alias: {
 svelte: path.resolve('node_modules', 'svelte')
 },
 extensions: ['.mjs', '.js', '.svelte'],
 mainFields: ['svelte', 'browser', 'module', 'main']
 }
};

Arrhythmia classification with stationary first order Markov process

Mon, 24 Dec 2018 00:00:00 +0000

Time series sequences can be described by their statistical properties like the mean level, trend, perodicity, autocorrelation, variability, and entropy. Most sequence classification models exploit these and other statistical differences between time series signals. However, some features are computationally expensive to calculate and they may not have sufficient discriminative power. Here I will show how Markov chains can be used to derive a simple discriminative statistical measure to separate time series sequences by summarizing the transition probabilities between consecutive timesteps.

A classic paper on ECG time series classification: A new method for detecting atrial fibrillation using R-R intervals by George Moody and Roger Mark (1990), uses markov chains to classify sequences of heart rate measurements as normal sinus rhythm (NS) and atrial fibrillation (AF). By using the sequence of heart rate values directly we don’t have to engineer new features and can work with the transition matrix derived from the observed sequence. This method is a good choice when signal variability is informative, as is the case in arrhythmia classification. The method is also computationally inexpensive and easy to deploy on low-power devices. The derived Markov score is also discriminative and can be used as a feature in downstream machine learning models.

The outline of the method is as follows: transitions between consecutive probabilities are summarized in a transition matrix assuming a stationary first-order Markov process (current value depends only on previous value). The log of the ratio of the two transition matrices for group A and group B yields the log-odds matrix $S$. The score is calculated as the sum of the log-odds for each new observed transition.

Transition matrix

In more detail:

Discretize heart rate into a number of bins. We’ll call these bins states.
Calculate a transition probability ($p_{i,j}$ = number of transitions from $state_i$ to $state_j$ / total number of transitions) between each state and store it in a transition matrix (above figure). We’ll call this matrix $T$. We want a transition matrix for each class. In our example we will be distinguishing a normal sinus rhythm ($T_\text{NS}$) from an arrhythmia like atrial fibrilation ($T_\text{AF}$).
Calculate the odds ratio of observing a state transiton in a normal rhythm and an abnormal rhythm: $p_{i,j}^{\text{NS}} / p_{i,j}^{\text{AF}}$. Store the odds ratios for each state transition in a score matrix $S$. In practice we can get the score matrix $S$ simply by dividing the two transition matrices: $S = T_\text{NS} / T_\text{AF}$.
Take the log of the score matrix so $S = \text{log}(T_\text{NS} / T_\text{AF})$. This gives us the log odds. We can interpret the score matrix as $S_{i,j}=0$ if there is an equal likelihood of observing a transition from $state_i$ to $state_j$ in both a normal sinus and AF rhythm. Note that taking the $log(A/B)$ is equivalent to subtracting $log(A)-log(B)$. It follows that positive values of $S$ ($S_{i,j}>0$) means a greater chance that the state transition from $state_i$ to $state_j$ comes from a normal rhythm and negative values of $S$ ($S_{i,j}<0$) means you are more likely to observe the state transition in an AF rhythm.
To get the final score we sum up the transition probabilities from the score matrix for each observation and choose an appropriate threshold at which we are confident the transition comes from an AF rhythm.

I will use ECG data from the Physionet 2017 Challenge to illustrate how a Markov process is used to classify time series signals. Participants in the challenge had to classify ECG signals collected from mobile sensors as normal sinus rhythm, atrial fibrillation, noise or other. For our purposes we ignore the noise and other labels and simply classify normal and atrial fibrillation.

Since the data is provided as raw ECG signals we first extract heart rate by detecting the R-peaks of the ECG beat and calculating the R to R intervals (heart rate = 60 / RR interval in seconds). The distribution of RR intervals between the two groups already shows a good separation between Normal and AF rhythms.

R-R interval distribution

Next, calculate the transition matrix for all RR intervals in the training data.

def transition_matrix(transitions):
 """Takes an array of discrete values and returns a transition matrix
 Args:
 transitions: list of states
 Returns:
 transition matrix, rows must sum to 1
 """
 n = 1 + max(transitions) # number of states
 M = np.zeros((n,n))
 for (i,j) in zip(transitions,transitions[1:]):
 M[i][j] += 1
 # now convert to probabilities:
 for row in M:
 s = sum(row)
 if s > 0:
 row[:] = [f/s for f in row]
 return np.array(M)

The rows of the transition matrix sum to 1 and shows the probability of observing a transition between RR intervals (in milliseconds) from time $t-1$ to $t$. The RR intervals were first discretized into bins of step size 50. The coarseness of the discretization is a hyperparameter that can be tuned to find the level that gives the best class separation.

Normal sinus rhythm (left), Atrial fibrillation (right) transition matrix

Dividing and taking the log of the two transition matrices gives the us the score matrix where each element is the log-odds ratio of observing a transition between two consecutive states. In the score matrix below, transitions in red are more likely to be observed in atrial fibrilation rhythms and transitions in green are more likely to be observed in normal sinus rhythm.

Score matrix

The final score is calculated by summing over transitions in a newly observed sequence of RR intervals.

$$M = \sum_{t=2}^{t=N} S(t-1,t)$$

The score can be plotted over time to show how the score changes with new observations. Examples for normal sinus rhythms have an increasing score and examples with atrial fibrillation have a decreasing score. We can use a cutoff threshold of zero such that a score > 0 will classify the example as a normal sinus rhythm.

Markov score

Some closing throughts: Why is the Markov score discriminative? Because we calculated the log-odds ratio between the two classes. The score works well in binary classification where the variability between consecutive observations is an informative discriminator. It may not work well in cases where the signal is periodic and the classification accuracy depends on information captured in the periodicity of the signal. In these cases a simple differencing or STL decomposition may help remove the periodic signal. The score presented here is also limited by the fact that it is first-order, meaning we look at only consecutive differences between time $t-1$ and $t$. In practice it may be better to use the score as a feature in a machine learning classifier since it captures a discriminiative attribute regarding the variability of the signal.

Nowcasting: Maintaining real time estimates of infrequently observed time series

Wed, 05 Dec 2018 00:00:00 +0000

Time series analysis appears in every disciple from physiology to retail pricing. A time series variable is typically measured sequentially at fixed intervals of time (often equispaced but not necessarily). Variables may be measured less frequently than theoretically possible for reasons of cost, effort, or convention. With local level linear trend models we can maintain realtime measures of infrequently measured values (see Predicting the Present with Bayesian Structural Time Series). The problem has been referred to as nowcasting because the goal is to maintain a current estimate of the value of a time series by forecasting the current value instead of the future value. The term itself is not very important as the task is essentially a standard forecasting problem.

Consider a measurement like US weekly initial claims for unemployment (ICNSA), which is a recession leading indicator. Can we learn this week’s number before it is released? To answer this question we would need a real time signal correlated with the outcome (ICNSA numbers). We can use Google Correlate to extract the top 100 search terms that are most correlated with the ICNSA signal. Google Correlate finds search terms that vary in a similar way to your own time series, ICNSA signal in our case. The 100 search term time series signals are our explanatory (also caled exogenous) variables that can be included as regressors to improve the ICNSA forecast performance. The idea is that contemporaneous signals (exogenous variables) are correlated in time with the unobserved signal (endogenous variable) we are trying to estimate and by regressing on these features can improve our forecast. The temporal structure in these observed signals can be exploited to infer the behaviour of an unobserved signal. Here we will explore using structural time series models that decompose a signal into additive components consisting a linear trend and a mean level.

US weekly initial claims for unemployment (ICNSA)

Brief description of structural time series models

The general approach to time series analysis is to first remove or model the parts that change through time to get a stationary series (a time series is stationary if its statistical properties, like variance, don’t change through time). Next, we use a time series model to capture the correlation in the stationary series. A series can be decomposed into:

trend components (long-term change in the mean level)
seasonality component (variation in mean that is periodic in nature and you generally know the period beforehand)
cycles (variation that oscillates but not according to some known or fixed period)
exogenous variables that have some correlation with the endogenous variable
noise

The various components can be combined additively to model the endogenous variable $y$ at time $t$. Such additive models are desirable because we can interpret each term, progressively increase model complexity, and easily diagnose model performance. More concretely a typical model will be written as:

$$ y_t=\mu_t+\gamma_t+\beta^Tx_t+\epsilon_t $$

where $y_t$ is the endogenous variable we want to forecast, $\mu_t$ captures changes in the mean level over time, $\gamma_t$ models the periodic nature of the signal, $\beta^T\boldsymbol{x}_t$ is a regression term with exogenous variables and $\epsilon_t$ is the noise term.

The local linear trend model decomposes the time series into a local level component and a trend component.

$$ \mu_t = \mu_{t-1}+\delta_{t-1}+u_t $$

$$ \delta_t = \delta_{t-1}+v_t $$

The current level of the trend is $\mu_t$, the current “slope” of the trend is $\delta_t$, and the noise terms are $u_t$ and $v_t$.

This kind of model is referred to as UnobservedComponents in statsmodels.

from statsmodels.tsa.statespace.structural import UnobservedComponents

# train on all time points before this and forecast time points after
interventionidx = 200

# df: dataframe with ICNSA and exogenous variables
# regression_columns: exogenous variables
intervention = df.index[interventionidx]
model = UnobservedComponents(
 df.loc[:intervention, 'ICNSA'].values,
 exog = df.loc[:intervention, regression_columns].values,
 level = 'local linear trend'
)
fit = model.fit(maxiter=1000)

We can compare a few models: without exogenous variables from Google Correlate, with the top 10 most correlated search terms, and with the bottom 10 least correlated search terms. The figures below show the ICNSA values in blue and the model predictions in red. The model is trained on observations until 2008 (vertical dashed line) and forecasts are made for the unobserved time after 2008. 95% confidence intervals are in grey.

ICNSA signal with model predictions

It’s clear that adding additional features to the model improves both the fit to the observed data and the forecast. But adding uncorrelated data can have undesired effects on your forecasts. The unobserved components model in statsmodels is unable to pick the best features since it does not have any kind of regularization. Ideally, we want to select only those correlated search terms that gives the best model fit and forecast. The original paper on Bayesian Structural Time Series model provides a methodology for feature selection.

In addition to applications in forecasting, state space models like the one described above can be used to infer the effect of an intervention, like an ad campaign, for counterfactual inference (see Inferring Causal Impact from Bayesian Structural Time-Series Models by Kay Brodersen et. al. (2015))

Useful references:

Original paper: Predicting the Present with Bayesian Structural Time Series
Inferring Causal Impact from Bayesian Structural Time-Series Models
Causal Impact R package
Rob Hyndmans books and papers, the teaching slides are particularly accessible

State space models and the Kalman filter

Tue, 02 Feb 2016 00:00:00 +0000

Linear state-space models are used in time-series analysis for filtering, prediction, and smoothing problems. They assume that the observations are generated linearly from a latent linear dynamical system. Although many real world processes are non-linear, the lineary makes the model easy to analyze and efficient to estimate. In addition, many non-linear systems can be approximated using linear models, thus the linear state-space model is an important tool for time-series applications.

Consider the basic structural model with a local level term and a trend term:

$$ y_t=\mu_t+\lambda_t w_t+\epsilon_t $$

$$ \mu_{t+1}=\mu_t+v_t+W_{1t} $$

$$ v_{t+1}=v_t+W_{2t} $$

$$ \lambda_{t+1}=\lambda_t+W_{3t} $$

where $\epsilon_t \sim N(0,\sigma_{y}^{2})$, $W_{1t} \sim N(0,\sigma_{\mu}^{2})$, and $W_{2t} \sim N(0,\sigma_{v}^{2})$. Here we allow the local level (intercept) and the trend (slope) to vary in time. Note the term local here is in contrast to global, where the level $\mu$ is fixed ($\sigma_{\mu}^{2}=0$) and there is a constant level across time.

In this case we have added an intervention variable $\lambda$ and $w$, where $\lambda$ is a weighting term and $w$ is a function where the value is zero before the intervention and unity after the intervention.

Setting all the noise terms $\eta=(\epsilon_t, W_{1t}, W_{2t})$ to zero yields the simple equation of a line with constant intercept and slope. At $t=1$:

$$ y_1=\mu_1 $$

$$ \mu_1=\mu_0+v_0 $$

$$ v_1=v_0 $$

$$ y_1=\mu_0 + v_0 $$

At $t=2$:

$$ y_2=\mu_2 $$

$$ \mu_2=\mu_1+v_1=\mu_0+v_0+v_0 $$

$$ v_2=v_1=v_0 $$

$$ y_2=\mu_0+2v_0 $$

At $t=3$:

$$y_3=\mu_2 $$

$$ \mu_3=\mu_2+v_2=\mu_0+v_0+v_0+b_0 $$

$$ v_3=v_2=v_1=v_0 $$

$$ y_3=\mu_0+3v_0 $$

Therefore, in this case the linear trend model simplifies to

$$ y_t=\mu_0+v_0g_t+\epsilon_t $$

where $g_t=t$ for $t=1,…,n$ is effectively time and $\mu_0$ and $v_0$ are the initial values of the level and the slope.

The state space model above can be expressed algebraically in one unified formulation. Using matrix algebra, these models can be written in the following general format:

$$ y_t=Z_{t}^{T}\alpha_t+\epsilon_t $$

$$ \alpha_{t+1}=T_t \alpha_t + R_t \eta_t $$

The first equation is the observation or measurement equation because it links the observed data with the unobserved latent state $\alpha$. The second equation is the transition or state equation because it defines how the latent state evolves over time. $\alpha$ is the state vector, $Z_t$ is the observation or design vector, $T_t$ is the transition matrix, $R_t$ is usually an identity matrix and in cases where it is not identity $R_t$ is called the selection matrix. Finally, $\eta$ is state disturbances.

We can express the local linear trend model in state space form:

$$ \alpha_t=\begin{pmatrix}\mu_t\v_t\end{pmatrix}, \quad \eta_t=\begin{pmatrix}\psi_t\\zeta_t\end{pmatrix}, \quad T_t=\begin{bmatrix}1 & 1\0 & 1\end{bmatrix}, \quad Z_t=\begin{pmatrix}1\0\end{pmatrix} $$

$$ Q_t=\begin{bmatrix}\sigma_{\mu}^2 & 0\0 & \sigma_{v}^2\end{bmatrix}, \quad R_t=\begin{bmatrix}1 & 0\0 & 1\end{bmatrix} $$

The primary tool for fitting state space model to data is the Kalman filter, which recursively computes the predictive distribution $p(\alpha_{t+1}\mid y_{1:t}))$ by combining $p(\alpha_{t}\mid y_{1:t-1}))$ with $y_t$ using a standard set of formulas that is logically equivalent to linear regression.

Intervention variables can be added to assess the influence of an external change or stimulus to the development in a time series. Three possible interventions are the level shift, slope shift, and a pulse where the value suddenly changes at the moment of the interventiona and than immediately returns to the value before the intervention took place. Changes in the value of level shift and slope shift are permanent after the intervention. A level shift can be expressed as follows:

$$ y_t=\mu_t+\lambda_t w_t+\epsilon_t $$

$$ \mu_{t+1}=\mu_t+W_{1t} $$

$$ \lambda_{t+1}=\lambda_t+W_{3t} $$

The dummy variable $w_t$ equals zero at all time points before the intervention and equals unity at the time points after the intervention.

The state space equations can also be cast as a probabilistic model such that for the measurement model we have $y_t\sim p(y_t \mid \alpha_t)$ and for the latent state model we have $\alpha_t\sim p(\alpha_t \mid \alpha_{t-1})$.

Resources:

An Introduction to State Space Time Series Analysis by Commandeur and Koopman

Pandoc static site generator

Wed, 20 May 2015 00:00:00 +0000

Pandoc is one of the most useful command-line document converters I’ve used. Entire websites can be generated from simple text files. I’m partial to the Markdown syntax, but really any plain-text file can be read into Pandoc and it will output HTML, PDF, and even Word Docs. This website for instance is built using a single Makefile that runs a Pandoc command.

pandoc -r markdown+simple_tables+table_captions+yaml_metadata_block+auto_identifiers
 +header_attributes+fenced_code_blocks+fenced_code_attributes+tex_math_dollars
 -w html --toc --mathjax --include-before-body=$(PANDOC)/navigation.html
 --template=$(PANDOC)/layout.html --css=$(PANDOC)/style.css -o $(ROOT)/$@ $<

Basically, for every Markdown .md file in the directory, Pandoc converts it to an HTML file styled using a given template. The Markdown file supports YAML headers and the variables are made available in the template file. Templates also support basic logic (if/else) and loops (for), which allows for some very smart template files that can change the HTML output based on the YAML headers in your Markdown file.

Pandoc also support syntax highlighting for embedded code.

Data science at the command line

Thu, 25 Dec 2014 00:00:00 +0000

The UNIX command line is a powerful tool for diving into large data files and piping specialized utilities to create summary statistics. This is a compilation of some useful commands.

Parsing CSV

Read the first line of a CSV file, which contains the column names, and list each column name in a new line after splitting by comma and stripping surrounding quotes. Notice awk can create an array and has operations like length that returns the number of elements in an array.

head -n 1 WebExtract.txt | awk '{ split($0, a, ","); max = length(a) }
 END { for (x=1; x<=max;x++) {gsub(/"/, "", a[x]); print a[x]} }'

Take the average of a column of integers in a CSV file after filtering for a category.

cat iris.csv | grep 'Iris-setosa' | awk -F "," '{ sum += $1; n += 1; }
 END { printf "%0.5f\n", sum/n }'

To select a single line in a CSV file and output one column.

sed -n '2405p;2405q' WebExtract.txt | awk -F "\",\"" '{ print $3; }'

We can also select a subset of lines in the middle of the document and return a single column.

sed -n '1000,1010p' WebExtract.txt | awk -F "\",\"" { 'print $3; }'

Count the number of lines with wc.

cat trip_data_1.csv | wc -l

JSON files can be parsed using jq. The following command parses a 35MB file for the attribute city, sorts the cities, and returns the number of occurrences.

cat ./yelp_train_academic_dataset_business.json | jq '.city' | sort | uniq -c

Parsing JSON

To test a command on a large file we can select just the first few lines.

head -n 100 ./yelp_train_academic_dataset_business.json | jq '.city' | sort | uniq -c

To further filter the output for occurrences greater than 5 we can use awk.

head -n 100 ./yelp_train_academic_dataset_business.json | jq '.city' |
 sort | uniq -c | awk '{if ($1 > 5) {print $0}}'

The above filter also works reasonably well on the entire dataset.

cat ./yelp_train_academic_dataset_business.json | jq '.city' |
 sort | uniq -c | awk '{if ($1 > 100) {print $0}}'

To get the total number of entires with over 100 occurrences we can take the sum.

cat ./yelp_train_academic_dataset_business.json | jq '.city' |
 sort | uniq -c | awk '{ if ($1 > 100) {sum += $1;} } END { print sum }'

To filter by a key within jq we can pipe commands and return only the city names.

head -n 3 ./yelp_train_academic_dataset_business.json | jq 'select(.city |
 contains("De Forest")) | .city'

Punchcard visualization using D3.js

Tue, 12 Aug 2014 00:00:00 +0000

Consider using a punchcard visualization when you want to represent both counts and proportionality over a time period across categorical data.

The graphic is particularly useful in interactive documents with hover effects showing the underlying count.

Data Mining PubMed

Wed, 12 Mar 2014 00:00:00 +0000

The National Institutes of Health provides a full programming interface to search PubMed called E-Utilities. Interacting with the PubMed database is conveniently through simple HTTP requests and returns the article metadata as XML. Every article in PubMed has a title, author, abstract, journal, year, volume, issue, pages, and keywords, amoung other metadata. Getting the metadata from PubMed, however, involves two separate queries. Very simply, the first query returns a list of PubMed IDs for articles matching the search criteria and the second query returns article data for a given PMID.

The workflow is divided into two parts:

Query E-Search passing it your search term and it returns a list of PMIDs that are used to query E-Fetch for the article metadata.

http://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi?term=electrical+stimulation
&amp;retmax=10&amp;tool=pmquery&amp;db=pubmed

E-Search returns a list of PMIDs:

 <eSearchResult>
 <Count>157380</Count>
 <RetMax>10</RetMax>
 <RetStart>0</RetStart>
 <IdList>
 <Id>23858010</Id>
 <Id>23856563</Id>
 <Id>23856146</Id>
 <Id>23855510</Id>
 <Id>23839460</Id>
 <Id>23839375</Id>
 <Id>23853340</Id>
 <Id>23853339</Id>
 <Id>23853324</Id>
 <Id>23853296</Id>
 </IdList>
 ...

Next, query E-Fetch for the article data. You can request multiple PMIDs at once and even the return type (XML, text, JSON). The API also supports pagination to iteratively get many thousands of results.

http://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi
?db=pubmed&id=23856563,23858010&retmode=xml

I’ve written several interfaces to access the PubMed API, including in PHP, Python, and C#. For instance, the Python script was written specifically to data-mine PubMed. Given a search term pmquery.py will query PubMed and save each article to a text file. For some search terms, like “transcranial magnetic stimulation” this results in over 9000 articles returned by Pubmed. So the process is iterative and can take some time (minutes). The PHP implementation provides a web-based search interface. For desktop based applications, see the C# code and the Scholared app for a working example.

Asif Rahman

ChartQL - A token-efficient SQL extension for generating Plotly charts

Installing Local Python Packages with uv tool

Disk-based caching with asset management in Python

Features

Basic Usage

Expiration Examples

Memoization Decorator

Asset Handling with Native Formats

Extensible Asset Handlers

Clearing Cache by Name or Function

Storage Backends

Eviction Policies

Pickle Fallback

Advanced Configuration

Implementation code

Document Chunking for LLMs

Algorithm Overview

Example: Berkshire Hathaway Annual Report

Implementation Code

Python Dependency Injection

Usage

Simple Dependency Injection

Async Dependencies

Dependency Caching

Disable Caching

Custom Fields

Dependency Overrides

Generator Dependencies (Context Managers)

Type Validation with Pydantic

Disable Type Casting

Implementation Code

Autostore: File Storage Made Simple

Why Use AutoStore?

Getting Started

Basic Usage

Cloud Storage (S3)

Supported Data Types

Configuration Options

S3StorageConfig

Advanced Features

Caching System

Custom Data Handlers

File Operations

Context Management

Multiple Storage Backends

Performance Considerations

Large File Handling

When to Use AutoStore

Comparison with Alternatives

Presskit: Database-driven static site generator

Key Features

Installation

Quick Start

Basic Usage

Writing Markdown Content

Creating HTML Templates

Configuration

Template Variables

Site Variables (site.*)

Build Variables (build.*)

Page Variables (page.*)

Data Variables (data.*)

Using Variables in Markdown

Data Sources and Queries

Configuring Data Sources

Adding Queries

Using Query Data in Templates

Page-Level Queries

Generating Pages

Generator Queries

Generator Configuration

Creating Generator Templates

Nested Queries

Commands

Build Commands

Development

Advanced Configuration

Full Configuration Example

Custom Filters

Site Variables (`site.*`)

Build Variables (`build.*`)

Page Variables (`page.*`)

Data Variables (`data.*`)