48 lines
2.4 KiB
Python
48 lines
2.4 KiB
Python
import asyncio
|
|
import csv
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
|
|
from assembly_dataset.config import Config
|
|
from assembly_dataset.pipeline import generate, export, validate
|
|
|
|
|
|
class CsvExportTests(unittest.TestCase):
|
|
def test_each_shard_has_readable_csv_with_same_records(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config = Config(output_dir=Path(tmp), target_examples=25, shard_size=2)
|
|
config.prompts.teacher_system = 'Правила'
|
|
config.prompts.training_system['full'] = 'Текст; с "кавычками"\nи новой строкой'
|
|
asyncio.run(generate(config))
|
|
export(config)
|
|
for path in Path(tmp).glob('sft/*/*/*.jsonl'):
|
|
counterpart = path.with_suffix('.csv')
|
|
self.assertTrue(counterpart.exists())
|
|
self.assertTrue(counterpart.read_bytes().startswith(b'\xef\xbb\xbf'))
|
|
originals = [json.loads(line) for line in path.read_text().splitlines()]
|
|
with counterpart.open(encoding='utf-8-sig', newline='') as file:
|
|
rows = list(csv.DictReader(file, delimiter=';'))
|
|
self.assertEqual(len(rows), len(originals))
|
|
for original, row in zip(originals, rows):
|
|
self.assertEqual(row['ID записи'], original['record_id'])
|
|
self.assertEqual(row['Системная инструкция'], original['messages'][0]['content'])
|
|
self.assertEqual(row['План'], original['messages'][2]['content'])
|
|
self.assertNotRegex(row['План'], r'part_\d|slot_\d')
|
|
if original['output_format'] == 'text':
|
|
self.assertNotIn('"skill"', row['План'])
|
|
self.assertTrue(row['План'])
|
|
self.assertEqual(validate(config)['examples_checked'], 25)
|
|
|
|
def test_corrupt_csv_is_detected(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config = Config(output_dir=Path(tmp), target_examples=1)
|
|
asyncio.run(generate(config))
|
|
export(config)
|
|
files = list(Path(tmp).glob('sft/*/*/*.csv'))
|
|
self.assertEqual(len(files), 2)
|
|
files[0].write_text('broken')
|
|
with self.assertRaisesRegex(ValueError, 'Повреждён'):
|
|
validate(config)
|