# josh-huang/ml-research-agent architecture

> Architecture diagram of the GitHub repository josh-huang/ml-research-agent, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/josh-huang/ml-research-agent
- Repository: https://github.com/josh-huang/ml-research-agent

## Overview

This repository is an autonomous recommender-model research agent: it loads KuaiRand data, establishes an FM baseline, iterates through evidence-based experiments, and produces a ranked test submission. The central loop is documented in the README; sampled source confirms the training, extended data loading, research retrieval, and submission paths. Some agent-tool wiring and evaluation details are documented but their implementation was not sampled, so uncertain internal links are omitted.

Main entry points: [`models/train.py`](https://github.com/josh-huang/ml-research-agent/blob/main/models/train.py) (Training loop), [`agent/main.py`](https://github.com/josh-huang/ml-research-agent/blob/main/agent/main.py) (Loop orchestrator), [`agent/react.py`](https://github.com/josh-huang/ml-research-agent/blob/main/agent/react.py) (ReAct researcher), [`agent/tools.py`](https://github.com/josh-huang/ml-research-agent/blob/main/agent/tools.py) (Experiment tools)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_agent["Research agent"]
  node_main["Loop orchestrator<br/>[main.py]"]
  node_react["ReAct researcher<br/>[react.py]"]
  node_llm["LLM client<br/>[llm.py]"]
  node_tools["Experiment tools<br/>[tools.py]"]
  node_research["Literature retrieval<br/>[research.py]"]
  node_eda["EDA summary<br/>[eda.py]"]
  node_search["Config search<br/>[search.py]"]
  node_state[("Run state<br/>[state.py]")]
  node_logger["Run logger<br/>[logger.py]"]
end

subgraph group_modeling["Modeling pipeline"]
  node_loader["Extended loader<br/>[data_loader.py]"]
  node_train["Training loop<br/>[train.py]"]
  node_fm["FM model<br/>[fm_torch.py]"]
  node_deepfm["DeepFM model<br/>[deepfm.py]"]
  node_din["DIN model<br/>[din.py]"]
  node_losses["Ranking losses<br/>[losses.py]"]
  node_evaluator["Official evaluator<br/>[evaluate.py]"]
end

subgraph group_delivery["Results and delivery"]
  node_submission["Submission generator<br/>[make_submission.py]"]
  node_submit["Submission writer<br/>[submit.py]"]
  node_csv["Ranked CSV<br/>[final.csv]"]
  node_logs[("Run records<br/>[run_log.jsonl]")]
  node_dashboard["Run dashboard<br/>[app.py]"]
end

node_operator(("Operator"))
node_anthropic(("Anthropic API"))
node_arxiv(("arXiv API"))
node_dataset[("KuaiRand data")]

node_operator -->|"starts run"| node_main
node_main -->|"loads data"| node_loader
node_dataset -->|"provides rows"| node_loader
node_main -->|"builds summary"| node_eda
node_main -->|"starts episode"| node_react
node_react -->|"requests reasoning"| node_llm
node_llm -.->|"sends prompts"| node_anthropic
node_react -->|"dispatches tools"| node_tools
node_react -->|"requests papers"| node_research
node_research -.->|"queries"| node_arxiv
node_react -->|"normalizes config"| node_search
node_main -->|"runs baseline"| node_tools
node_tools -->|"runs experiment"| node_train
node_train -->|"builds model"| node_fm
node_train -->|"builds model"| node_deepfm
node_train -->|"builds model"| node_din
node_train -->|"computes objective"| node_losses
node_train -->|"scores predictions"| node_evaluator
node_tools -->|"updates state"| node_state
node_tools -->|"records run"| node_logger
node_logger -->|"appends records"| node_logs
node_submission -->|"loads features"| node_loader
node_submission -->|"trains best config"| node_train
node_submission -->|"writes scores"| node_submit
node_submit -->|"creates file"| node_csv
node_dashboard -->|"reads trail"| node_logs
node_dashboard -->|"reads state"| node_state

click node_main "https://github.com/josh-huang/ml-research-agent/blob/main/agent/main.py"
click node_react "https://github.com/josh-huang/ml-research-agent/blob/main/agent/react.py"
click node_llm "https://github.com/josh-huang/ml-research-agent/blob/main/agent/llm.py"
click node_tools "https://github.com/josh-huang/ml-research-agent/blob/main/agent/tools.py"
click node_research "https://github.com/josh-huang/ml-research-agent/blob/main/agent/research.py"
click node_eda "https://github.com/josh-huang/ml-research-agent/blob/main/agent/eda.py"
click node_search "https://github.com/josh-huang/ml-research-agent/blob/main/agent/search.py"
click node_state "https://github.com/josh-huang/ml-research-agent/blob/main/agent/state.py"
click node_logger "https://github.com/josh-huang/ml-research-agent/blob/main/agent/logger.py"
click node_loader "https://github.com/josh-huang/ml-research-agent/blob/main/models/data_loader.py"
click node_train "https://github.com/josh-huang/ml-research-agent/blob/main/models/train.py"
click node_fm "https://github.com/josh-huang/ml-research-agent/blob/main/models/fm_torch.py"
click node_deepfm "https://github.com/josh-huang/ml-research-agent/blob/main/models/deepfm.py"
click node_din "https://github.com/josh-huang/ml-research-agent/blob/main/models/din.py"
click node_losses "https://github.com/josh-huang/ml-research-agent/blob/main/models/losses.py"
click node_evaluator "https://github.com/josh-huang/ml-research-agent/blob/main/kuairand-starter-kit/kuairand-starter-kit/evaluate.py"
click node_submission "https://github.com/josh-huang/ml-research-agent/blob/main/submission/make_submission.py"
click node_submit "https://github.com/josh-huang/ml-research-agent/blob/main/kuairand-starter-kit/kuairand-starter-kit/submit.py"
click node_csv "https://github.com/josh-huang/ml-research-agent/blob/main/submission/final.csv"
click node_logs "https://github.com/josh-huang/ml-research-agent/blob/main/run_logs/run_log.jsonl"
click node_dashboard "https://github.com/josh-huang/ml-research-agent/blob/main/dashboard/app.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_main,node_react,node_llm,node_tools,node_research,node_eda,node_search,node_state,node_logger toneBlue
class node_loader,node_train,node_fm,node_deepfm,node_din,node_losses,node_evaluator,node_dataset toneAmber
class node_submission,node_submit,node_csv,node_logs,node_dashboard,node_anthropic,node_arxiv toneMint
class node_operator toneIndigo
```

## Components

### Research agent

- [Loop orchestrator](https://github.com/josh-huang/ml-research-agent/blob/main/agent/main.py): `agent/main.py`
- [ReAct researcher](https://github.com/josh-huang/ml-research-agent/blob/main/agent/react.py): `agent/react.py`
- [LLM client](https://github.com/josh-huang/ml-research-agent/blob/main/agent/llm.py): `agent/llm.py`
- [Experiment tools](https://github.com/josh-huang/ml-research-agent/blob/main/agent/tools.py): `agent/tools.py`
- [Literature retrieval](https://github.com/josh-huang/ml-research-agent/blob/main/agent/research.py): `agent/research.py`
- [EDA summary](https://github.com/josh-huang/ml-research-agent/blob/main/agent/eda.py): `agent/eda.py`
- [Config search](https://github.com/josh-huang/ml-research-agent/blob/main/agent/search.py): `agent/search.py`
- [Run state](https://github.com/josh-huang/ml-research-agent/blob/main/agent/state.py): `agent/state.py`
- [Run logger](https://github.com/josh-huang/ml-research-agent/blob/main/agent/logger.py): `agent/logger.py`

### Modeling pipeline

- [Extended loader](https://github.com/josh-huang/ml-research-agent/blob/main/models/data_loader.py): `models/data_loader.py`
- [Training loop](https://github.com/josh-huang/ml-research-agent/blob/main/models/train.py): `models/train.py`
- [FM model](https://github.com/josh-huang/ml-research-agent/blob/main/models/fm_torch.py): `models/fm_torch.py`
- [DeepFM model](https://github.com/josh-huang/ml-research-agent/blob/main/models/deepfm.py): `models/deepfm.py`
- [DIN model](https://github.com/josh-huang/ml-research-agent/blob/main/models/din.py): `models/din.py`
- [Ranking losses](https://github.com/josh-huang/ml-research-agent/blob/main/models/losses.py): `models/losses.py`
- [Official evaluator](https://github.com/josh-huang/ml-research-agent/blob/main/kuairand-starter-kit/kuairand-starter-kit/evaluate.py): `kuairand-starter-kit/kuairand-starter-kit/evaluate.py`

### Results and delivery

- [Submission generator](https://github.com/josh-huang/ml-research-agent/blob/main/submission/make_submission.py): `submission/make_submission.py`
- [Submission writer](https://github.com/josh-huang/ml-research-agent/blob/main/kuairand-starter-kit/kuairand-starter-kit/submit.py): `kuairand-starter-kit/kuairand-starter-kit/submit.py`
- [Ranked CSV](https://github.com/josh-huang/ml-research-agent/blob/main/submission/final.csv): `submission/final.csv`
- [Run records](https://github.com/josh-huang/ml-research-agent/blob/main/run_logs/run_log.jsonl): `run_logs/run_log.jsonl`
- [Run dashboard](https://github.com/josh-huang/ml-research-agent/blob/main/dashboard/app.py): `dashboard/app.py`

### Other components

- **Operator**
- **Anthropic API**
- **arXiv API**
- **KuaiRand data**

## Connections

- Operator → Loop orchestrator: starts run
- Loop orchestrator → Extended loader: loads data
- KuaiRand data → Extended loader: provides rows
- Loop orchestrator → EDA summary: builds summary
- Loop orchestrator → ReAct researcher: starts episode
- ReAct researcher → LLM client: requests reasoning
- LLM client → Anthropic API: sends prompts
- ReAct researcher → Experiment tools: dispatches tools
- ReAct researcher → Literature retrieval: requests papers
- Literature retrieval → arXiv API: queries
- ReAct researcher → Config search: normalizes config
- Loop orchestrator → Experiment tools: runs baseline
- Experiment tools → Training loop: runs experiment
- Training loop → FM model: builds model
- Training loop → DeepFM model: builds model
- Training loop → DIN model: builds model
- Training loop → Ranking losses: computes objective
- Training loop → Official evaluator: scores predictions
- Experiment tools → Run state: updates state
- Experiment tools → Run logger: records run
- Run logger → Run records: appends records
- Submission generator → Extended loader: loads features
- Submission generator → Training loop: trains best config
- Submission generator → Submission writer: writes scores
- Submission writer → Ranked CSV: creates file
- Run dashboard → Run records: reads trail
- Run dashboard → Run state: reads state