| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import yaml |
|
|
| from .config import TrainingConfig |
| from .data.prepare import load_configured_datasets, save_prepared |
| from .data.split import split_dataset |
| from .evaluate import compare_models, evaluate_model, save_result |
| from .train_baseline import train_tfidf |
|
|
|
|
| def _config(path: Path) -> TrainingConfig: |
| raw = yaml.safe_load(path.read_text(encoding="utf-8")) or {} |
| cfg = TrainingConfig.from_mapping(raw.get("training", {})) |
| datasets = (yaml.safe_load(path.read_text(encoding="utf-8")) or {}).get("datasets", []) |
| return TrainingConfig.from_env(TrainingConfig.from_mapping({**cfg.__dict__, "datasets": datasets})) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Expense Tracker ML training pipeline") |
| sub = parser.add_subparsers(dest="command", required=True) |
| prep = sub.add_parser("prepare") |
| prep.add_argument("--config", type=Path, default=Path("config/datasets.yaml")) |
| prep.add_argument("--output", type=Path, default=Path("data/prepared/transactions.parquet")) |
| train = sub.add_parser("train") |
| train.add_argument("--config", type=Path, default=Path("config/datasets.yaml")) |
| train.add_argument("--model", choices=["tfidf", "transformer", "both"], default="both") |
| args = parser.parse_args() |
|
|
| if args.command == "prepare": |
| cfg = _config(args.config) |
| frame = load_configured_datasets(cfg.datasets, cfg.data_dir) |
| print(json.dumps(save_prepared(frame, args.output), indent=2)) |
| return |
|
|
| cfg = _config(args.config) |
| prepared = cfg.data_dir / "prepared/transactions.parquet" |
| frame = __import__("pandas").read_parquet(prepared) |
| splits = split_dataset(frame, cfg.seed, cfg.test_size, cfg.validation_size) |
| results = [] |
| if args.model in {"tfidf", "both"}: |
| model = train_tfidf(splits.train, cfg) |
| results.append(evaluate_model(model, splits.test, "tfidf")) |
| if len(splits.india_holdout): |
| save_result(evaluate_model(model, splits.india_holdout, "tfidf-india"), cfg.artifacts_dir / "tfidf-india.json") |
| if args.model in {"transformer", "both"}: |
| from .train_transformer import train_transformer |
| model = train_transformer(splits.train, splits.validation, cfg) |
| results.append(evaluate_model(model, splits.test, "transformer")) |
| if len(splits.india_holdout): |
| save_result(evaluate_model(model, splits.india_holdout, "transformer-india"), cfg.artifacts_dir / "transformer-india.json") |
| comparison = compare_models(results) |
| (cfg.artifacts_dir / "comparison.json").write_text(json.dumps(comparison, indent=2), encoding="utf-8") |
| print(json.dumps(comparison, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
