File size: 735 Bytes
7a60a87 | 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 | import unittest
from pathlib import Path
from sglang.utils import execute_shell_command
CACHE_DIR = Path(__file__).parent.parent.parent.joinpath("cache")
class TestPrepareData(unittest.TestCase):
def test_prepare_sharegpt(self):
sharegpt_train_path = CACHE_DIR.joinpath("dataset", "sharegpt_train.jsonl")
if sharegpt_train_path.exists():
# delete the file
sharegpt_train_path.unlink()
process = execute_shell_command(
"python scripts/prepare_data.py --dataset sharegpt"
)
process.wait()
self.assertEqual(process.returncode, 0)
self.assertTrue(sharegpt_train_path.exists())
if __name__ == "__main__":
unittest.main(verbosity=2)
|