<commit_before>from xformer import data
from xformer.data import record_file
from xformer.scripts import split_pile

import pytest
import os.path
import tempfile

from dataclasses import dataclass
from typing import List, Tuple


@pytest.skip("deprecated due to needing data")
def test_pile():
    ds = data.PileDataset(
        os.path.join(os.path.dirname(__file__), "../data/pile/train/00.jsonl.zst"),
        n_ctx=1024,
    )
    it = iter(ds)
    for _ in range(10):
        ex = next(it)
        assert ex.ndim == 1
        assert ex.shape[0] <= 1024


def test_record_file():
    with tempfile.TemporaryFile() as fh:
        w = os.fdopen(os.dup(fh.fileno()), "w+b")

        with record_file.Writer(w) as wf:
            wf.write(b"hello, world")
            wf.write(b"goodnight, moon")

        fh.seek(0)
<commit_msg>whoops<commit_after>from xformer import data
from xformer.data import record_file
from xformer.scripts import split_pile

import pytest
import os.path
import tempfile

from dataclasses import dataclass
from typing import List, Tuple


@pytest.mark.skip("deprecated due to needing data")
def test_pile():
    ds = data.PileDataset(
        os.path.join(os.path.dirname(__file__), "../data/pile/train/00.jsonl.zst"),
        n_ctx=1024,
    )
    it = iter(ds)
    for _ in range(10):
        ex = next(it)
        assert ex.ndim == 1
        assert ex.shape[0] <= 1024


def test_record_file():
    with tempfile.TemporaryFile() as fh:
        w = os.fdopen(os.dup(fh.fileno()), "w+b")

        with record_file.Writer(w) as wf:
            wf.write(b"hello, world")
            wf.write(b"goodnight, moon")

        fh.seek(0)

{'fraction_non_alphanumeric': 0.12016908212560387, 'fraction_numerical': 0.018115942028985508, 'mean_word_length': 2.731981981981982, 'pattern_counts': {'":': 0, '<': 5, '<?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': '3709685', 'n_tokens_mistral': 604, 'n_tokens_neox': 567, 'n_words': 140}