2023-02-08 18:48:41 +00:00
|
|
|
import typing as t
|
|
|
|
|
import contextlib
|
|
|
|
|
import dataclasses
|
|
|
|
|
import logging
|
|
|
|
|
import pymongo
|
|
|
|
|
|
|
|
|
|
from .config import TRAINING_SET_SIZE, VALIDATION_SET_SIZE, EVALUATION_SET_SIZE
|
2023-02-12 04:11:58 +00:00
|
|
|
from .database import SampleFunc, CachedDatasetFunc, store_cache, load_cache, delete_cache
|
2023-02-08 18:48:41 +00:00
|
|
|
|
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@dataclasses.dataclass
|
|
|
|
|
class Caches:
|
|
|
|
|
"""
|
|
|
|
|
Container for the three generators that can create datasets.
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
training: CachedDatasetFunc
|
|
|
|
|
validation: CachedDatasetFunc
|
|
|
|
|
evaluation: CachedDatasetFunc
|
|
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
|
@contextlib.contextmanager
|
|
|
|
|
def from_database_samples(cls, collection: pymongo.collection.Collection, sample_func: SampleFunc) -> t.ContextManager["Caches"]:
|
|
|
|
|
"""
|
|
|
|
|
Create a new caches object from a database collection and a sampling function.
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
log.debug("Gathering datasets...")
|
|
|
|
|
reviews_training = sample_func(collection, TRAINING_SET_SIZE.__wrapped__)
|
|
|
|
|
reviews_validation = sample_func(collection, VALIDATION_SET_SIZE.__wrapped__)
|
|
|
|
|
reviews_evaluation = sample_func(collection, EVALUATION_SET_SIZE.__wrapped__)
|
|
|
|
|
|
|
|
|
|
log.debug("Caching datasets...")
|
|
|
|
|
store_cache(reviews_training, "./data/training")
|
|
|
|
|
store_cache(reviews_validation, "./data/validation")
|
|
|
|
|
store_cache(reviews_evaluation, "./data/evaluation")
|
|
|
|
|
|
|
|
|
|
log.debug("Loading dataset caches...")
|
|
|
|
|
training_cache = load_cache("./data/training")
|
|
|
|
|
validation_cache = load_cache("./data/validation")
|
|
|
|
|
evaluation_cache = load_cache("./data/evaluation")
|
|
|
|
|
|
|
|
|
|
yield Caches(training=training_cache, validation=validation_cache, evaluation=evaluation_cache)
|
|
|
|
|
|
|
|
|
|
log.debug("Cleaning up caches...")
|
|
|
|
|
delete_cache("./data/training")
|
|
|
|
|
delete_cache("./data/validation")
|
|
|
|
|
delete_cache("./data/evaluation")
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def ensure_clean():
|
|
|
|
|
log.debug("Ensuring there are no leftover caches...")
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
delete_cache("./data/training")
|
|
|
|
|
except FileNotFoundError:
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
delete_cache("./data/validation")
|
|
|
|
|
except FileNotFoundError:
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
delete_cache("./data/evaluation")
|
|
|
|
|
except FileNotFoundError:
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
__all__ = (
|
|
|
|
|
"Caches",
|
|
|
|
|
)
|
