41 lines
1.0 KiB
Python
41 lines
1.0 KiB
Python
|
"""
|
||
|
|
Step 1: Generate companion conversation dataset using LLM.
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
python scripts/generate_data.py --config configs/data_generation.yaml
|
||
|
|
"""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import yaml
|
||
|
|
from pathlib import Path
|
||
|
|
from src.data.data_generator import ConversationGenerator
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--config", default="configs/data_generation.yaml")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
with open(args.config) as f:
|
||
|
|
cfg = yaml.safe_load(f)
|
||
|
|
|
||
|
|
Path(cfg["output"]["raw_dir"]).mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
generator = ConversationGenerator(
|
||
|
|
api_type=cfg["api"]["type"],
|
||
|
|
model=cfg["api"]["model"],
|
||
|
|
)
|
||
|
|
|
||
|
|
count = generator.generate_dataset(
|
||
|
|
output_path=cfg["output"]["output_file"],
|
||
|
|
total_samples=cfg["generation"]["total_samples"],
|
||
|
|
samples_per_category=cfg["generation"]["samples_per_category"],
|
||
|
|
delay=cfg["generation"]["delay"],
|
||
|
|
)
|
||
|
|
|
||
|
|
print(f"Generated {count} samples → {cfg['output']['output_file']}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|