# kiyas-mahmud/webagent architecture

> Architecture diagram of the GitHub repository kiyas-mahmud/webagent, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-27). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/kiyas-mahmud/webagent
- Repository: https://github.com/kiyas-mahmud/webagent

## Overview

1. **What this repo is**
   - This is an offline ML research/training repo for a **failure-aware web agent**.
   - The central design is **backbone-agnostic**: one model template stays fixed, while only the front-end encoder/backbone changes across many variants.
   - The core prediction stack always ends in **one fused 768-d embedding** feeding **four task heads** for the four “pillars” of the agent: detect, diagnose, recover, remember.

2. **Top-level system boundary**
   - The repository is organized around a single pipeline:
     **dataset preparation → verification → dataloading → model forward pass → multi-task loss → training/checkpointing → evaluation/reporting**.
   - There is no web serving layer here; the architecture is for **training and evaluation**, not deployment.

3. **Data subsystem**
   - `scripts/setup_data.py` and `src/web_agent/data/verify_dataset.py` are the entry points for turning the raw dataset into a usable local training corpus.
   - The dataset is built around **Mind2Web-style step data** with failure augmentation, plus a shared `images/` folder that should not be split.
   - `src/web_agent/data/dataset.py`, `gold_dataset.py`, and `dataloader.py` define the main sample format and batching logic, including the **five mode filters** mentioned in the README.
   - `gold_dataset.py` / `gold_dataloader.py` appear to support the “gold” subset used for the highest-quality training path.

4. **Model boundary: fixed core, swappable encoder**
   - The `src/web_agent/models/` package is the main architectural boundary.
   - It is split into:
     - `encoders/` for backbone-specific input processing,
     - `fusion.py` for combining modalities,
     - `adapter.py` for projecting backbone outputs into the fixed hidden size,
     - `heads.py` for the four task heads,
     - `loss.py` for the shared weighted objective,
     - `model.py` for assembling the full end-to-end network.
   - This is the key diagram separation: **backbone varies, downstream task stack does not**.

5. **Two encoder families**
   - The README describes two major patterns:
     - **Dual-encoder models**: `vision_enc + text_enc -> cross-attention -> 768-d`
     - **Unified VLM models**: `VLM(image+text) -> linear adapter -> 768-d`
   - In code, these correspond to different implementations under `src/web_agent/models/encoders/vision.py`, `text.py`, and `vlm.py`.
   - This is a meaningful architectural split because it changes where multimodal fusion happens: either **before** the shared 768-d representation or **inside** the backbone itself.

6. **Task heads and loss**
   - The model always produces one shared latent embedding, then branches into **four fixed heads**.
   - `heads.py` likely maps that shared representation into the four outputs required by the project’s pillars.
   - `loss.py` combines them into a **single weighted multi-task loss**, which the README says is fixed across all backbones.
   - `labels.py` is the likely source of the canonical target definitions used by both training and evaluation.

7. **Memory subsystem**
   - `src/web_agent/memory/index.py` is a distinct subsystem for the **remember** pillar.
   - Architecturally, this suggests an external retrieval/index component rather than a pure parametric prediction head.
   - It should be drawn as a sidecar capability that consumes past interaction traces or embeddings and can influence the memory-related output path.

8. **Training pipeline**
   - `src/web_agent/train/` is staged:
     - `smoke_test.py` for tiny correctness runs,
     - `mini_train.py` for a small but real training slice,
     - `train.py` and `trainer.py` for the full training loop.
   - The README emphasizes a strict progression: **SMOKE → MINI → FULL → EVAL → CHECKPOINT**.
   - This makes the training pipeline a gated system, not a single monolithic “train” command.

9. **Evaluation subsystem**
   - `src/web_agent/eval/metrics.py` and `evaluate.py` implement model assessment.
   - The README indicates evaluation spans **three test splits**, so the evaluator is not a single-score endpoint.
   - `results/results_table.csv` and the `results/` folder act as the persistent reporting layer for comparison across backbones and runs.

10. **Baseline models**
   - `src/web_agent/baselines/` contains standalone reference implementations:
     `b1_random.py`, `b2_mindact.py`, `b3_clip_mlp.py`, `b4_layoutlmv3.py`.
   - These are important in the architecture diagram as **comparative branches**, not part of the main training stack.
   - They provide performance anchors against the shared four-head architecture.

11. **Configuration and orchestration**
   - `configs/base.yaml` and `configs/backbones/*.yaml` define the model variants.
   - The config layer is the main mechanism for selecting which encoder/backbone is plugged into the shared model core.
   - `scripts/run_model.py` is the main orchestration entry point, with stage selection for smoke/mini/full/eval.
   - `scripts/run_gold.py` appears to be a specialized path for the gold subset.

12. **Runtime and environment assumptions**
   - The project is designed to run locally or on notebook/GPU environments, with the README calling out **Kaggle T4** and **Vast.ai A100** as compute targets.
   - `notebooks/kaggle_train.ipynb` is intentionally thin: it installs the package and calls the same script entry point, so notebooks are wrappers rather than a separate code path.

13. **Cross-cutting utilities**
   - `src/web_agent/utils/` provides shared infrastructure: checkpointing, class weighting, logging, and seeding.
   - These are supporting services used across training/eval rather than user-facing features.
   - `src/web_agent/config.py` likely normalizes config loading and run parameters across scripts.

14. **Testing and verification**
   - `tests/test_shapes.py` shows the repo uses at least structural tests for tensor shapes and likely finite-loss sanity.
   - Combined with the smoke stage and dataset verification, the repository has a strong “verify early” architecture: data sanity, shape sanity, then scaled training.

15. **Docs as architecture source of truth**
   - The `docs/` and `learn/` directories are unusually important here: they contain the project spec, architecture design, implementation plan, and phased training notes.
   - For diagramming, treat these as the authoritative design layer that explains why the code is organized into fixed heads + swappable backbones + staged training.

Main entry points: [`scripts/run_model.py`](https://github.com/kiyas-mahmud/webagent/blob/Code/scripts/run_model.py) (Run model), [`src/web_agent/eval/evaluate.py`](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/eval/evaluate.py) (Evaluate), [`src/web_agent/train/trainer.py`](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/train/trainer.py) (Trainer), [`src/web_agent/data/dataloader.py`](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/data/dataloader.py) (Batching)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_g_data["Data"]
  node_n_setup["Setup data<br/>script<br/>[setup_data.py]"]
  node_n_verify["Verify data<br/>validator<br/>[verify_dataset.py]"]
  node_n_dataset["Dataset<br/>sample schema<br/>[dataset.py]"]
  node_n_gold_dataset["Gold set<br/>gold dataset<br/>[gold_dataset.py]"]
  node_n_dataloader["Batching<br/>dataloader<br/>[dataloader.py]"]
  node_n_gold_loader["Gold loader<br/>gold dataloader<br/>[gold_dataloader.py]"]
end

subgraph group_g_model["Model"]
  node_n_text_enc["Text enc<br/>encoder<br/>[text.py]"]
  node_n_vision_enc["Vision enc<br/>encoder<br/>[vision.py]"]
  node_n_vlm_enc["VLM enc<br/>encoder<br/>[vlm.py]"]
  node_n_fusion{{"Fusion<br/>[fusion.py]"}}
  node_n_adapter["Adapter<br/>projection<br/>[adapter.py]"]
  node_n_heads["Task heads<br/>4-head output<br/>[heads.py]"]
  node_n_loss["Multi-task loss<br/>weighted loss<br/>[loss.py]"]
  node_n_model["Model core<br/>assembled model<br/>[model.py]"]
  node_n_memory[("Memory index<br/>retrieval sidecar<br/>[index.py]")]
end

subgraph group_g_train["Training"]
  node_n_smoke["Smoke test<br/>train stage<br/>[smoke_test.py]"]
  node_n_mini["Mini train<br/>train stage<br/>[mini_train.py]"]
  node_n_trainer["Trainer<br/>loop<br/>[trainer.py]"]
  node_n_train["Full train<br/>train entry<br/>[train.py]"]
end

subgraph group_g_eval["Evaluation"]
  node_n_eval["Evaluate<br/>eval entry<br/>[evaluate.py]"]
  node_n_metrics["Metrics<br/>metric suite<br/>[metrics.py]"]
  node_n_results["Results<br/>reporting<br/>[results_table.csv]"]
end

subgraph group_g_ops["Ops"]
  node_n_run_model["Run model<br/>orchestrator<br/>[run_model.py]"]
  node_n_run_gold["Run gold<br/>orchestrator<br/>[run_gold.py]"]
  node_n_config["Config<br/>[config.py]"]
  node_n_utils["Utilities<br/>shared infra"]
end

subgraph group_g_refs["References"]
  node_n_baselines["Baselines<br/>reference models"]
  node_n_docs["Docs<br/>spec layer"]
end

node_n_setup -->|"prepare corpus"| node_n_verify
node_n_verify -->|"validate"| node_n_dataset
node_n_dataset -->|"batch"| node_n_dataloader
node_n_dataset -->|"subset"| node_n_gold_dataset
node_n_gold_dataset -->|"batch"| node_n_gold_loader
node_n_dataloader -->|"feed"| node_n_smoke
node_n_dataloader -->|"feed"| node_n_mini
node_n_dataloader -->|"feed"| node_n_trainer
node_n_gold_loader -->|"feed"| node_n_train
node_n_text_enc -->|"encode"| node_n_fusion
node_n_vision_enc -->|"encode"| node_n_fusion
node_n_vlm_enc -->|"encode"| node_n_adapter
node_n_fusion -->|"combine"| node_n_adapter
node_n_adapter -->|"768-d latent"| node_n_heads
node_n_heads -->|"predict"| node_n_loss
node_n_memory -.->|"remember signal"| node_n_heads
node_n_model -->|"train step"| node_n_loss
node_n_smoke -->|"reuse loop"| node_n_trainer
node_n_mini -->|"reuse loop"| node_n_trainer
node_n_train -->|"run loop"| node_n_trainer
node_n_run_model -->|"stage"| node_n_smoke
node_n_run_model -->|"stage"| node_n_mini
node_n_run_model -->|"stage"| node_n_train
node_n_run_model -->|"stage"| node_n_eval
node_n_run_gold -->|"gold path"| node_n_gold_loader
node_n_config -->|"select backbone"| node_n_run_model
node_n_config -->|"build model"| node_n_model
node_n_utils -.->|"support"| node_n_trainer
node_n_utils -.->|"support"| node_n_eval
node_n_model -->|"predict"| node_n_eval
node_n_eval -->|"score"| node_n_metrics
node_n_eval -->|"report"| node_n_results
node_n_baselines -.->|"compare"| node_n_results
node_n_docs -.->|"specify"| node_n_config
node_n_docs -.->|"design"| node_n_model

click node_n_setup "https://github.com/kiyas-mahmud/webagent/blob/Code/scripts/setup_data.py"
click node_n_verify "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/data/verify_dataset.py"
click node_n_dataset "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/data/dataset.py"
click node_n_gold_dataset "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/data/gold_dataset.py"
click node_n_dataloader "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/data/dataloader.py"
click node_n_gold_loader "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/data/gold_dataloader.py"
click node_n_text_enc "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/encoders/text.py"
click node_n_vision_enc "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/encoders/vision.py"
click node_n_vlm_enc "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/encoders/vlm.py"
click node_n_fusion "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/fusion.py"
click node_n_adapter "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/adapter.py"
click node_n_heads "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/heads.py"
click node_n_loss "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/loss.py"
click node_n_model "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/model.py"
click node_n_memory "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/memory/index.py"
click node_n_smoke "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/train/smoke_test.py"
click node_n_mini "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/train/mini_train.py"
click node_n_trainer "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/train/trainer.py"
click node_n_train "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/train/train.py"
click node_n_eval "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/eval/evaluate.py"
click node_n_metrics "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/eval/metrics.py"
click node_n_run_model "https://github.com/kiyas-mahmud/webagent/blob/Code/scripts/run_model.py"
click node_n_run_gold "https://github.com/kiyas-mahmud/webagent/blob/Code/scripts/run_gold.py"
click node_n_config "https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/config.py"
click node_n_utils "https://github.com/kiyas-mahmud/webagent/tree/Code/src/web_agent/utils"
click node_n_baselines "https://github.com/kiyas-mahmud/webagent/tree/Code/src/web_agent/baselines"
click node_n_docs "https://github.com/kiyas-mahmud/webagent/blob/Code/docs/ARCHITECTURE_DESIGN.md"
click node_n_results "https://github.com/kiyas-mahmud/webagent/blob/Code/results/results_table.csv"

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_n_setup,node_n_verify,node_n_dataset,node_n_gold_dataset,node_n_dataloader,node_n_gold_loader toneBlue
class node_n_text_enc,node_n_vision_enc,node_n_vlm_enc,node_n_fusion,node_n_adapter,node_n_heads,node_n_loss,node_n_model,node_n_memory toneAmber
class node_n_smoke,node_n_mini,node_n_trainer,node_n_train toneMint
class node_n_eval,node_n_metrics,node_n_results toneRose
class node_n_run_model,node_n_run_gold,node_n_config,node_n_utils toneIndigo
class node_n_baselines,node_n_docs toneTeal
```

## Components

### Data

Dataset preparation, verification, loading, and gold subset handling.

- [Setup data](https://github.com/kiyas-mahmud/webagent/blob/Code/scripts/setup_data.py): script · Builds the local training corpus from raw data and images. · `scripts/setup_data.py`
- [Verify data](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/data/verify_dataset.py): validator · Checks dataset integrity and readiness before training. · `src/web_agent/data/verify_dataset.py`
- [Dataset](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/data/dataset.py): sample schema · Main dataset format for step-level web-agent examples. · `src/web_agent/data/dataset.py`
- [Gold set](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/data/gold_dataset.py): gold dataset · Higher-quality subset used for the gold training path. · `src/web_agent/data/gold_dataset.py`
- [Batching](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/data/dataloader.py): dataloader · Standard batching and mode filtering for training samples. · `src/web_agent/data/dataloader.py`
- [Gold loader](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/data/gold_dataloader.py): gold dataloader · Loader tuned for the gold subset training path. · `src/web_agent/data/gold_dataloader.py`

### Model

Backbone-agnostic multimodal model stack with fixed heads and shared loss.

- [Text enc](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/encoders/text.py): encoder · Text-side encoder used in dual-encoder backbones. · `src/web_agent/models/encoders/text.py`
- [Vision enc](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/encoders/vision.py): encoder · Image-side encoder used in dual-encoder backbones. · `src/web_agent/models/encoders/vision.py`
- [VLM enc](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/encoders/vlm.py): encoder · Unified multimodal backbone for image-plus-text inputs. · `src/web_agent/models/encoders/vlm.py`
- [Fusion](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/fusion.py): multimodal fusion · Combines modalities or backbone outputs before the fixed latent. · `src/web_agent/models/fusion.py`
- [Adapter](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/adapter.py): projection · Projects backbone outputs into the shared 768-d representation. · `src/web_agent/models/adapter.py`
- [Task heads](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/heads.py): 4-head output · Four fixed outputs for detect, diagnose, recover, and remember. · `src/web_agent/models/heads.py`
- [Multi-task loss](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/loss.py): weighted loss · Shared weighted objective across all backbones and heads. · `src/web_agent/models/loss.py`
- [Model core](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/models/model.py): assembled model · End-to-end network wiring encoders, fusion, adapter, heads, and loss. · `src/web_agent/models/model.py`
- [Memory index](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/memory/index.py): retrieval sidecar · External memory subsystem for the remember pillar. · `src/web_agent/memory/index.py`

### Training

Staged training entry points and trainer loop.

- [Smoke test](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/train/smoke_test.py): train stage · Tiny correctness run for early validation. · `src/web_agent/train/smoke_test.py`
- [Mini train](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/train/mini_train.py): train stage · Small but real training slice before full runs. · `src/web_agent/train/mini_train.py`
- [Trainer](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/train/trainer.py): loop · Shared training loop and checkpoint coordination. · `src/web_agent/train/trainer.py`
- [Full train](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/train/train.py): train entry · Main full-scale training entry point. · `src/web_agent/train/train.py`

### Evaluation

Metrics and evaluation reporting.

- [Evaluate](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/eval/evaluate.py): eval entry · Runs evaluation over the test splits. · `src/web_agent/eval/evaluate.py`
- [Metrics](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/eval/metrics.py): metric suite · Computes task and split-level evaluation metrics. · `src/web_agent/eval/metrics.py`
- [Results](https://github.com/kiyas-mahmud/webagent/blob/Code/results/results_table.csv): reporting · Persistent run comparison table and output artifacts. · `results/results_table.csv`

### Ops

Configuration, utilities, and orchestration.

- [Run model](https://github.com/kiyas-mahmud/webagent/blob/Code/scripts/run_model.py): orchestrator · Main stage selector for smoke, mini, full, and eval flows. · `scripts/run_model.py`
- [Run gold](https://github.com/kiyas-mahmud/webagent/blob/Code/scripts/run_gold.py): orchestrator · Specialized orchestration path for the gold subset. · `scripts/run_gold.py`
- [Config](https://github.com/kiyas-mahmud/webagent/blob/Code/src/web_agent/config.py): config layer · Loads base and backbone-specific run configuration. · `src/web_agent/config.py`
- [Utilities](https://github.com/kiyas-mahmud/webagent/tree/Code/src/web_agent/utils): shared infra · Checkpointing, class weights, logging, and seeding support. · `src/web_agent/utils`

### References

Baseline comparisons and project documentation.

- [Baselines](https://github.com/kiyas-mahmud/webagent/tree/Code/src/web_agent/baselines): reference models · Standalone comparison models for performance anchoring. · `src/web_agent/baselines`
- [Docs](https://github.com/kiyas-mahmud/webagent/blob/Code/docs/ARCHITECTURE_DESIGN.md): spec layer · Architecture, training plan, and implementation guidance. · `docs/ARCHITECTURE_DESIGN.md`

## Connections

- Setup data → Verify data: prepare corpus · Raw data is turned into a local dataset and then verified.
- Verify data → Dataset: validate · Verification produces a ready dataset schema.
- Dataset → Batching: batch · Standard samples are batched for training.
- Dataset → Gold set: subset · A higher-quality gold subset is derived from the same corpus.
- Gold set → Gold loader: batch · Gold samples use a dedicated loader path.
- Batching → Smoke test: feed · Standard batches drive smoke runs.
- Batching → Mini train: feed · Standard batches drive mini training.
- Batching → Trainer: feed · Standard batches feed the full trainer loop.
- Gold loader → Full train: feed · Gold runs can route through the full training entry.
- Text enc → Fusion: encode · Text features enter the multimodal fusion stage.
- Vision enc → Fusion: encode · Vision features enter the multimodal fusion stage.
- VLM enc → Adapter: encode · Unified VLM backbones produce outputs for projection.
- Fusion → Adapter: combine · Fused features are projected into the fixed latent size.
- Adapter → Task heads: 768-d latent · Shared embedding fans out to the four task heads.
- Task heads → Multi-task loss: predict · Head outputs are combined under a weighted objective.
- Memory index → Task heads: remember signal · External memory supports the remember-related output path.
- Model core → Multi-task loss: train step · The assembled model computes the shared multitask loss.
- Smoke test → Trainer: reuse loop · Smoke runs exercise the same trainer machinery.
- Mini train → Trainer: reuse loop · Mini training reuses the shared trainer machinery.
- Full train → Trainer: run loop · Full training is executed through the trainer.
- Run model → Smoke test: stage · The main runner can dispatch smoke runs.
- Run model → Mini train: stage · The main runner can dispatch mini runs.
- Run model → Full train: stage · The main runner can dispatch full training.
- Run model → Evaluate: stage · The main runner can dispatch evaluation.
- Run gold → Gold loader: gold path · Specialized orchestration targets the gold data path.
- Config → Run model: select backbone · Config chooses which backbone variant is assembled.
- Config → Model core: build model · Configuration determines the instantiated model stack.
- Utilities → Trainer: support · Shared utilities support training and checkpointing.
- Utilities → Evaluate: support · Shared utilities also support evaluation runs.
- Model core → Evaluate: predict · Trained models are evaluated on test splits.
- Evaluate → Metrics: score · Evaluation delegates scoring to the metrics layer.
- Evaluate → Results: report · Evaluation writes comparison artifacts to results.
- Baselines → Results: compare · Baseline runs provide performance anchors in reporting.
- Docs → Config: specify · Documentation describes the configuration and staged workflow.
- Docs → Model core: design · Docs define the fixed-head, swappable-backbone architecture.