<commit_before># Lint as: python3
"""Metrics for evaluating models."""

import json
import numpy as np
import xarray as xr


def train_test_split_time(data, split_day):
  """Split data into training_data (before split_day) and test_data (after).

  Args:
    data: an xr.Dataset containing the ground-truth infections.
    split_day: a time.coord at which to split the data into train and test sets.

  Returns:
    train_data: an xr.Dataset containing the training data for a model.
    test_data: an xr.Dataset containing the testing data for a model.
  """
  # everything before split_day
  train_data = data.sel(time=slice(None, split_day - 1)).copy()
  # everything after split_day
  test_data = data.sel(time=slice(split_day, None)).copy()
  # drop all variables that won't exist on real test data
  things_included = ['location', 'time']
  if hasattr(test_data, 'dynamic_covariates'):
    things_included.append('dynamic_covariates')
  test_data = test_data[things_included]

  train_data.attrs['split'] = 'train'
  train_data.attrs['split_spec'] = 'train_test_split_time'
  test_data.attrs['split'] = 'test'
  test_data.attrs['split_spec'] = 'train_test_split_time'
<commit_msg>Google export<commit_after># Lint as: python3
"""Metrics for evaluating models."""

import json
import numpy as np
import xarray as xr


def train_test_split_time(data, split_day):
  """Split data into training_data (before split_day) and test_data (after).

  Args:
    data: an xr.Dataset containing the ground-truth infections.
    split_day: a time.coord at which to split the data into train and test sets.

  Returns:
    train_data: an xr.Dataset containing the training data for a model.
    test_data: an xr.Dataset containing the testing data for a model.
  """
  # everything before split_day
  train_data = data.sel(time=slice(None, split_day - 1)).copy()
  # everything after split_day
  test_data = data.sel(time=slice(split_day, None)).copy()
  # drop all variables that won't exist on real test data
  things_included = ['location', 'time']
  if hasattr(test_data, 'dynamic_covariates'):
    things_included.append('dynamic_covariates')
  if 'original_time' in test_data.data_vars.keys():
    things_included.append('original_time')
  test_data = test_data[things_included]

  train_data.attrs['split'] = 'train'
  train_data.attrs['split_spec'] = 'train_test_split_time'
  test_data.attrs['split'] = 'test'
  test_data.attrs['split_spec'] = 'train_test_split_time'

{'fraction_non_alphanumeric': 0.10453808752025932, 'fraction_numerical': 0.0016207455429497568, 'mean_word_length': 4.795774647887324, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7242143', 'n_tokens_mistral': 800, 'n_tokens_neox': 778, 'n_words': 232}