Spaces:
Runtime error
Runtime error
File size: 14,064 Bytes
e4f9cbe 55dc3dd e4f9cbe e9a1c18 e4f9cbe e9a1c18 e4f9cbe e9a1c18 e4f9cbe e9a1c18 e4f9cbe e9a1c18 e4f9cbe e9a1c18 e4f9cbe e9a1c18 e4f9cbe e9a1c18 e4f9cbe e9a1c18 e4f9cbe e9a1c18 e4f9cbe e9a1c18 e4f9cbe 55dc3dd e4f9cbe e9a1c18 e4f9cbe e9a1c18 e4f9cbe e9a1c18 e4f9cbe e9a1c18 cf614fd e9a1c18 cf614fd e9a1c18 cf614fd e9a1c18 cf614fd e9a1c18 cf614fd e9a1c18 cf614fd e9a1c18 cf614fd e9a1c18 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 |
"""Utilities for working with datasets."""
import json
import math
import os
import pickle
import pprint
import secrets
from collections.abc import Iterable
from typing import Any, Callable, Generator, Iterator, Optional, Sequence, TypeVar, Union, cast
import numpy as np
import pyarrow as pa
from ..parquet_writer import ParquetWriter
from ..schema import (
PATH_WILDCARD,
TEXT_SPAN_END_FEATURE,
TEXT_SPAN_START_FEATURE,
UUID_COLUMN,
VALUE_KEY,
Field,
Item,
PathTuple,
Schema,
VectorKey,
field,
schema,
schema_to_arrow_schema,
)
from ..signals.signal import EMBEDDING_KEY, Signal
from ..utils import file_exists, log, open_file
_KEYS_SUFFIX = '.keys.pkl'
_EMBEDDINGS_SUFFIX = '.npy'
def is_primitive(obj: object) -> bool:
"""Returns True if the object is a primitive."""
if isinstance(obj, (str, bytes, np.ndarray, int, float)):
return True
if isinstance(obj, Iterable):
return False
return True
def _replace_embeddings_with_none(input: Union[Item, Item]) -> Union[Item, Item]:
if isinstance(input, np.ndarray):
return None
if isinstance(input, dict):
return {k: _replace_embeddings_with_none(v) for k, v in input.items()}
if isinstance(input, list):
return [_replace_embeddings_with_none(v) for v in input]
return input
def replace_embeddings_with_none(input: Union[Item, Item]) -> Item:
"""Replaces all embeddings with None."""
return cast(Item, _replace_embeddings_with_none(input))
def lilac_span(start: int, end: int, metadata: dict[str, Any] = {}) -> Item:
"""Creates a lilac span item, representing a pointer to a slice of text."""
return {VALUE_KEY: {TEXT_SPAN_START_FEATURE: start, TEXT_SPAN_END_FEATURE: end}, **metadata}
def lilac_embedding(start: int, end: int, embedding: Optional[np.ndarray]) -> Item:
"""Creates a lilac embedding item, representing a vector with a pointer to a slice of text."""
return lilac_span(start, end, {EMBEDDING_KEY: embedding})
Tflatten = TypeVar('Tflatten', object, np.ndarray)
def _flatten(input: Union[Iterator, object], is_primitive_predicate: Callable[[object],
bool]) -> Generator:
"""Flattens a nested iterable."""
if is_primitive_predicate(input):
yield input
elif isinstance(input, dict):
yield input
elif is_primitive(input):
yield input
else:
for elem in cast(Iterator, input):
yield from _flatten(elem, is_primitive_predicate)
def flatten(input: Union[Iterator, Iterable, Tflatten],
is_primitive_predicate: Callable[[object], bool] = is_primitive) -> Iterator[Tflatten]:
"""Flattens a nested iterator.
Primitives and dictionaries are not flattened. The user can also provide a predicate to determine
what is a primitive.
"""
return _flatten(input, is_primitive_predicate)
def count_primitives(input: Union[Iterable, Iterator]) -> int:
"""Iterate through each element of the input, flattening each one, computing a count.
Sum the final set of counts. This is the important iterable not to exhaust.
"""
return sum((len(list(flatten(i))) for i in input))
def _wrap_value_in_dict(input: Union[object, dict], props: PathTuple) -> Union[object, dict]:
# If the signal produced no value, or nan, we should return None so the parquet value is sparse.
if isinstance(input, float) and math.isnan(input):
input = None
for prop in reversed(props):
input = {prop: input}
return input
def _unflatten(flat_input: Iterator[list[object]],
original_input: Union[Iterable, object]) -> Union[list, dict]:
"""Unflattens a flattened iterable according to the original iterable's structure."""
if is_primitive(original_input):
return next(flat_input)
else:
values: Iterable
if isinstance(original_input, dict):
values = original_input.values()
else:
values = cast(Iterable, original_input)
return [_unflatten(flat_input, orig_elem) for orig_elem in values]
def unflatten(flat_input: Union[Iterable, Iterator], original_input: Union[Iterable,
object]) -> list:
"""Unflattens a flattened iterable according to the original iterable's structure."""
return cast(list, _unflatten(iter(flat_input), original_input))
def _wrap_in_dicts(input: Union[object, Iterable[object]],
spec: list[PathTuple]) -> Union[object, Iterable[object]]:
"""Wraps an object or iterable in a dict according to the spec."""
props = spec[0] if spec else tuple()
if len(spec) == 1:
return _wrap_value_in_dict(input, props)
if input is None or isinstance(input, float) and math.isnan(input):
# Return empty dict for missing inputs.
return {}
res = [_wrap_in_dicts(elem, spec[1:]) for elem in cast(Iterable, input)]
return _wrap_value_in_dict(res, props)
def wrap_in_dicts(input: Iterable[object], spec: list[PathTuple]) -> Iterable[object]:
"""Wraps an object or iterable in a dict according to the spec."""
return [_wrap_in_dicts(elem, spec) for elem in input]
def _merge_field_into(schema: Field, destination: Field) -> None:
if isinstance(schema, Field):
destination.signal = destination.signal or schema.signal
destination.dtype = destination.dtype or schema.dtype
if schema.fields:
destination.fields = destination.fields or {}
for field_name, subfield in schema.fields.items():
if field_name not in destination.fields:
destination.fields[field_name] = subfield.copy(deep=True)
else:
_merge_field_into(subfield, destination.fields[field_name])
elif schema.repeated_field:
if not destination.repeated_field:
raise ValueError('Failed to merge schemas. Origin schema is repeated, but destination is not')
_merge_field_into(schema.repeated_field, destination.repeated_field)
else:
if destination.dtype != schema.dtype:
raise ValueError(f'Failed to merge schemas. Origin schema has dtype {schema.dtype}, '
f'but destination has dtype {destination.dtype}')
def merge_schemas(schemas: Sequence[Union[Schema, Field]]) -> Schema:
"""Merge a list of schemas."""
merged_schema = Schema(fields={})
for s in schemas:
_merge_field_into(cast(Field, s), cast(Field, merged_schema))
return merged_schema
def schema_contains_path(schema: Schema, path: PathTuple) -> bool:
"""Check if a schema contains a path."""
current_field = cast(Field, schema)
for path_part in path:
# When we reach a value key, the schema should have a dtype defined on it. If not, it is derived
# and this schema does not contain the value.
if path_part == VALUE_KEY:
return current_field.dtype is not None
if path_part == PATH_WILDCARD:
if current_field.repeated_field is None:
return False
current_field = current_field.repeated_field
else:
if current_field.fields is None or path_part not in current_field.fields:
return False
current_field = current_field.fields[str(path_part)]
return True
def create_signal_schema(signal: Signal, source_path: PathTuple, current_schema: Schema) -> Schema:
"""Create a schema describing the enriched fields added an enrichment."""
leafs = current_schema.leafs
# Validate that the enrich fields are actually a valid leaf path.
if source_path not in leafs:
raise ValueError(f'"{source_path}" is not a valid leaf path. Leaf paths: {leafs.keys()}')
signal_schema = signal.fields()
signal_schema.signal = signal.dict()
enriched_schema = field(fields={signal.key(is_computed_signal=True): signal_schema})
for path_part in reversed(source_path):
if path_part == PATH_WILDCARD:
enriched_schema = Field(repeated_field=enriched_schema)
else:
enriched_schema = Field(fields={path_part: enriched_schema})
if not enriched_schema.fields:
raise ValueError('This should not happen since enriched_schema always has fields (see above)')
return schema({UUID_COLUMN: 'string', **cast(dict, enriched_schema.fields)})
def write_item_embeddings_to_disk(keys: Iterable[str], embeddings: Iterable[object],
output_dir: str, shard_index: int, num_shards: int) -> str:
"""Write a set of embeddings to disk."""
output_path_prefix = embedding_index_filename_prefix(output_dir, shard_index, num_shards)
# Restrict the keys to only those that are embeddings.
def embedding_predicate(input: Any) -> bool:
return isinstance(input, np.ndarray)
flat_keys = flatten_keys(keys, embeddings, is_primitive_predicate=embedding_predicate)
flat_embeddings = flatten(embeddings, is_primitive_predicate=embedding_predicate)
embedding_vectors: list[np.ndarray] = []
embedding_keys: list[VectorKey] = []
for key, lilac_embedding in zip(flat_keys, flat_embeddings):
if not key or not lilac_embedding or EMBEDDING_KEY not in lilac_embedding:
# Sparse embeddings may not have an embedding for every key.
continue
# We use squeeze here because embedding functions can return outer dimensions of 1.
embedding_vectors.append(lilac_embedding[EMBEDDING_KEY].reshape(-1))
embedding_keys.append(key)
embedding_vectors = np.array(embedding_vectors)
# Write the embedding index and the ordered UUID column to disk so they can be joined later.
with open_file(output_path_prefix + _EMBEDDINGS_SUFFIX, 'wb') as f:
np.save(f, embedding_vectors, allow_pickle=False)
with open_file(output_path_prefix + _KEYS_SUFFIX, 'wb') as f:
pickle.dump(embedding_keys, f)
return output_path_prefix
def read_embedding_index(filepath_prefix: str) -> tuple[list[VectorKey], np.ndarray]:
"""Reads the embedding index for a column from disk."""
if not file_exists(filepath_prefix + _EMBEDDINGS_SUFFIX):
raise ValueError(F'Embedding index does not exist at path {filepath_prefix}. '
'Please run dataset.compute_signal() on the embedding signal first.')
# Read the embedding index from disk.
embeddings = np.load(filepath_prefix + _EMBEDDINGS_SUFFIX, allow_pickle=False)
with open_file(filepath_prefix + _KEYS_SUFFIX, 'rb') as f:
index_keys: list[VectorKey] = pickle.load(f)
return index_keys, embeddings
def write_items_to_parquet(items: Iterable[Item], output_dir: str, schema: Schema,
filename_prefix: str, shard_index: int,
num_shards: int) -> tuple[str, int]:
"""Write a set of items to a parquet file, in columnar format."""
arrow_schema = schema_to_arrow_schema(schema)
out_filename = parquet_filename(filename_prefix, shard_index, num_shards)
filepath = os.path.join(output_dir, out_filename)
f = open_file(filepath, mode='wb')
writer = ParquetWriter(schema)
writer.open(f)
num_items = 0
for item in items:
# Add a UUID column.
if UUID_COLUMN not in item:
item[UUID_COLUMN] = secrets.token_urlsafe(nbytes=12) # 16 base64 characters.
if os.getenv('DEBUG'):
try:
_validate(item, arrow_schema)
except Exception as e:
raise ValueError(f'Error validating item: {json.dumps(item)}') from e
writer.write(item)
num_items += 1
writer.close()
f.close()
return out_filename, num_items
def _validate(item: Item, schema: pa.Schema) -> None:
# Try to parse the item using the inferred schema.
try:
pa.RecordBatch.from_pylist([item], schema=schema)
except pa.ArrowTypeError:
log('Failed to parse arrow item using the arrow schema.')
log('Item:')
log(pprint.pformat(item, indent=2))
log('Arrow schema:')
log(schema)
raise # Re-raise the same exception, same stacktrace.
def parquet_filename(prefix: str, shard_index: int, num_shards: int) -> str:
"""Return the filename for a parquet file."""
return f'{prefix}-{shard_index:05d}-of-{num_shards:05d}.parquet'
def _flatten_keys(uuid: str, nested_input: Iterable, location: list[int],
is_primitive_predicate: Callable[[object], bool]) -> Iterator[VectorKey]:
if is_primitive_predicate(nested_input) or is_primitive(nested_input) or isinstance(
nested_input, dict):
yield (uuid, *location)
return
for i, input in enumerate(nested_input):
yield from _flatten_keys(uuid, input, [*location, i], is_primitive_predicate)
def flatten_keys(
uuids: Iterable[str],
nested_input: Iterable,
is_primitive_predicate: Callable[[object],
bool] = is_primitive) -> Iterator[Optional[VectorKey]]:
"""Flatten the uuid keys of a nested input."""
for uuid, input in zip(uuids, nested_input):
if input is None:
yield None
continue
yield from _flatten_keys(uuid, input, [], is_primitive_predicate)
def embedding_index_filename_prefix(output_dir: str, shard_index: int, num_shards: int) -> str:
"""Return the filename prefix for the embedding index."""
npy_filename = f'embeddings-{shard_index:05d}-of-{num_shards:05d}'
return os.path.join(output_dir, npy_filename)
Tin = TypeVar('Tin')
Tout = TypeVar('Tout')
def sparse_to_dense_compute(
sparse_input: Iterator[Optional[Tin]],
func: Callable[[Iterable[Tin]], Iterable[Tout]]) -> Iterator[Optional[Tout]]:
"""Densifies the input before calling the provided `func` and sparsifies the output."""
locations: list[int] = []
total_size: int = 0
def densify(x: Iterator[Optional[Tin]]) -> Iterator[Tin]:
nonlocal locations, total_size
for i, value in enumerate(x):
total_size += 1
if value is not None:
locations.append(i)
yield value
dense_input = densify(sparse_input)
dense_output = iter(func(dense_input))
index = 0
location_index = 0
while True:
try:
out = next(dense_output)
out_index = locations[location_index]
while index < out_index:
yield None
index += 1
yield out
location_index += 1
index += 1
except StopIteration:
while index < total_size:
yield None
index += 1
return
|