# lingyi0110/mai2bert architecture

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

- Interactive diagram: https://gitdiagram.com/lingyi0110/mai2bert
- Repository: https://github.com/lingyi0110/mai2bert

## Overview

mai2bert is a self-supervised pretraining and supervised chart-difficulty prediction framework for structured maimai event sequences. Its main lifecycle is to discover and parse chart files, build augmented event representations and splits, store processed data, train or fine-tune sequence models, then load a checkpoint to predict difficulty and uncertainty for chart inputs. The graph preserves both training and inference branches. The README is empty, and several modules were not sampled; their paths identify subsystems, but uncertain internal call links are omitted.

Main entry points: [`src/preprocessing/prepare.py`](https://github.com/lingyi0110/mai2bert/blob/main/src/preprocessing/prepare.py) (Preparation pipeline), [`src/models/regression.py`](https://github.com/lingyi0110/mai2bert/blob/main/src/models/regression.py) (Difficulty regressor), [`src/inference/batch.py`](https://github.com/lingyi0110/mai2bert/blob/main/src/inference/batch.py) (Batch prediction), [`train.py`](https://github.com/lingyi0110/mai2bert/blob/main/train.py) (Training command)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_prep["Chart preparation"]
  node_prepare_cli["Prepare command<br/>[prepare.py]"]
  node_discovery["Chart discovery<br/>[discovery.py]"]
  node_simai_parser["Simai parser<br/>[simai.py]"]
  node_ma2_parser["MA2 parser<br/>[ma2.py]"]
  node_representation["Event representation<br/>[representation.py]"]
  node_augmentation["Chart augmentation"]
  node_split["Dataset splits<br/>[split.py]"]
  node_prepare_pipeline["Preparation pipeline<br/>[prepare.py]"]
end

subgraph group_data["Dataset pipeline"]
  node_store[("Processed HDF5<br/>[store.py]")]
  node_dataset["Event dataset<br/>[dataset.py]"]
  node_batching["Batch collation<br/>[batching.py]"]
  node_masking["Pretraining masking<br/>[masking.py]"]
end

subgraph group_model["Sequence models"]
  node_embedding["Event embedding<br/>[embedding.py]"]
  node_encoder["Sequence encoder<br/>[encoder.py]"]
  node_pretraining_model["Pretraining model<br/>[pretraining.py]"]
  node_regression_model["Difficulty regressor<br/>[regression.py]"]
end

subgraph group_training["Training"]
  node_train_cli["Training command<br/>[train.py]"]
  node_pretrain["Pretraining loop<br/>[pretrain.py]"]
  node_finetune["Regression training<br/>[regression.py]"]
  node_transfer["Transfer training<br/>[transfer.py]"]
end

subgraph group_inference["Prediction"]
  node_inference_cli["Inference command<br/>[infer.py]"]
  node_batch_inference["Batch prediction<br/>[batch.py]"]
  node_inference_loader["Checkpoint loader<br/>[loader.py]"]
  node_predict["Prediction runtime<br/>[predict.py]"]
  node_prediction_output["Prediction files"]
end

node_chart_inputs(("Chart files"))
node_checkpoint["Model checkpoint"]

node_chart_inputs -->|"supplies charts"| node_prepare_cli
node_prepare_cli -->|"runs preparation"| node_prepare_pipeline
node_prepare_pipeline -->|"discovers charts"| node_discovery
node_prepare_pipeline -->|"parses charts"| node_simai_parser
node_prepare_pipeline -->|"builds events"| node_representation
node_prepare_pipeline -->|"creates variants"| node_augmentation
node_prepare_pipeline -->|"assigns splits"| node_split
node_prepare_pipeline -->|"writes arrays"| node_store
node_chart_inputs -->|"supplies charts"| node_inference_cli
node_inference_cli -->|"dispatches batch"| node_batch_inference
node_batch_inference -->|"parses Simai"| node_simai_parser
node_batch_inference -->|"parses MA2"| node_ma2_parser
node_batch_inference -->|"submits items"| node_predict
node_inference_cli -->|"loads model"| node_inference_loader
node_checkpoint -->|"provides weights"| node_inference_loader
node_inference_loader -->|"restores model"| node_regression_model
node_predict -->|"predicts difficulty"| node_regression_model
node_batch_inference -->|"writes CSV and JSON"| node_prediction_output
node_train_cli -->|"starts pretraining"| node_pretrain
node_train_cli -->|"starts regression"| node_finetune
node_train_cli -->|"starts transfer"| node_transfer
node_store -->|"supplies events"| node_dataset
node_dataset -->|"provides samples"| node_batching
node_batching -->|"collates for masking"| node_masking
node_masking -->|"provides masked batches"| node_pretraining_model
node_batching -->|"provides batches"| node_regression_model
node_pretraining_model -->|"encodes events"| node_embedding
node_regression_model -->|"encodes events"| node_embedding
node_embedding -->|"builds representations"| node_encoder
node_pretrain -->|"trains model"| node_pretraining_model
node_finetune -->|"trains regressor"| node_regression_model
node_transfer -->|"adapts model"| node_regression_model
node_finetune -->|"saves weights"| node_checkpoint

click node_prepare_cli "https://github.com/lingyi0110/mai2bert/blob/main/prepare.py"
click node_discovery "https://github.com/lingyi0110/mai2bert/blob/main/src/preprocessing/discovery.py"
click node_simai_parser "https://github.com/lingyi0110/mai2bert/blob/main/src/lib/parser/simai.py"
click node_ma2_parser "https://github.com/lingyi0110/mai2bert/blob/main/src/lib/parser/ma2.py"
click node_representation "https://github.com/lingyi0110/mai2bert/blob/main/src/preprocessing/representation.py"
click node_augmentation "https://github.com/lingyi0110/mai2bert/tree/main/src/lib/augmentation"
click node_split "https://github.com/lingyi0110/mai2bert/blob/main/src/preprocessing/split.py"
click node_prepare_pipeline "https://github.com/lingyi0110/mai2bert/blob/main/src/preprocessing/prepare.py"
click node_store "https://github.com/lingyi0110/mai2bert/blob/main/src/preprocessing/store.py"
click node_dataset "https://github.com/lingyi0110/mai2bert/blob/main/src/data/dataset.py"
click node_batching "https://github.com/lingyi0110/mai2bert/blob/main/src/data/batching.py"
click node_masking "https://github.com/lingyi0110/mai2bert/blob/main/src/data/masking.py"
click node_embedding "https://github.com/lingyi0110/mai2bert/blob/main/src/models/embedding.py"
click node_encoder "https://github.com/lingyi0110/mai2bert/blob/main/src/models/encoder.py"
click node_pretraining_model "https://github.com/lingyi0110/mai2bert/blob/main/src/models/pretraining.py"
click node_regression_model "https://github.com/lingyi0110/mai2bert/blob/main/src/models/regression.py"
click node_train_cli "https://github.com/lingyi0110/mai2bert/blob/main/train.py"
click node_pretrain "https://github.com/lingyi0110/mai2bert/blob/main/src/training/pretrain.py"
click node_finetune "https://github.com/lingyi0110/mai2bert/blob/main/src/training/regression.py"
click node_transfer "https://github.com/lingyi0110/mai2bert/blob/main/src/training/transfer.py"
click node_inference_cli "https://github.com/lingyi0110/mai2bert/blob/main/infer.py"
click node_batch_inference "https://github.com/lingyi0110/mai2bert/blob/main/src/inference/batch.py"
click node_inference_loader "https://github.com/lingyi0110/mai2bert/blob/main/src/inference/loader.py"
click node_predict "https://github.com/lingyi0110/mai2bert/blob/main/src/inference/predict.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_prepare_cli,node_discovery,node_simai_parser,node_ma2_parser,node_representation,node_augmentation,node_split,node_prepare_pipeline toneBlue
class node_store,node_dataset,node_batching,node_masking toneAmber
class node_embedding,node_encoder,node_pretraining_model,node_regression_model toneMint
class node_train_cli,node_pretrain,node_finetune,node_transfer toneRose
class node_inference_cli,node_batch_inference,node_inference_loader,node_predict,node_prediction_output,node_chart_inputs,node_checkpoint toneIndigo
```

## Components

### Chart preparation

- [Prepare command](https://github.com/lingyi0110/mai2bert/blob/main/prepare.py): `prepare.py`
- [Chart discovery](https://github.com/lingyi0110/mai2bert/blob/main/src/preprocessing/discovery.py): `src/preprocessing/discovery.py`
- [Simai parser](https://github.com/lingyi0110/mai2bert/blob/main/src/lib/parser/simai.py): `src/lib/parser/simai.py`
- [MA2 parser](https://github.com/lingyi0110/mai2bert/blob/main/src/lib/parser/ma2.py): `src/lib/parser/ma2.py`
- [Event representation](https://github.com/lingyi0110/mai2bert/blob/main/src/preprocessing/representation.py): `src/preprocessing/representation.py`
- [Chart augmentation](https://github.com/lingyi0110/mai2bert/tree/main/src/lib/augmentation): `src/lib/augmentation`
- [Dataset splits](https://github.com/lingyi0110/mai2bert/blob/main/src/preprocessing/split.py): `src/preprocessing/split.py`
- [Preparation pipeline](https://github.com/lingyi0110/mai2bert/blob/main/src/preprocessing/prepare.py): `src/preprocessing/prepare.py`

### Dataset pipeline

- [Processed HDF5](https://github.com/lingyi0110/mai2bert/blob/main/src/preprocessing/store.py): `src/preprocessing/store.py`
- [Event dataset](https://github.com/lingyi0110/mai2bert/blob/main/src/data/dataset.py): `src/data/dataset.py`
- [Batch collation](https://github.com/lingyi0110/mai2bert/blob/main/src/data/batching.py): `src/data/batching.py`
- [Pretraining masking](https://github.com/lingyi0110/mai2bert/blob/main/src/data/masking.py): `src/data/masking.py`

### Sequence models

- [Event embedding](https://github.com/lingyi0110/mai2bert/blob/main/src/models/embedding.py): `src/models/embedding.py`
- [Sequence encoder](https://github.com/lingyi0110/mai2bert/blob/main/src/models/encoder.py): `src/models/encoder.py`
- [Pretraining model](https://github.com/lingyi0110/mai2bert/blob/main/src/models/pretraining.py): `src/models/pretraining.py`
- [Difficulty regressor](https://github.com/lingyi0110/mai2bert/blob/main/src/models/regression.py): `src/models/regression.py`

### Training

- [Training command](https://github.com/lingyi0110/mai2bert/blob/main/train.py): `train.py`
- [Pretraining loop](https://github.com/lingyi0110/mai2bert/blob/main/src/training/pretrain.py): `src/training/pretrain.py`
- [Regression training](https://github.com/lingyi0110/mai2bert/blob/main/src/training/regression.py): `src/training/regression.py`
- [Transfer training](https://github.com/lingyi0110/mai2bert/blob/main/src/training/transfer.py): `src/training/transfer.py`

### Prediction

- [Inference command](https://github.com/lingyi0110/mai2bert/blob/main/infer.py): `infer.py`
- [Batch prediction](https://github.com/lingyi0110/mai2bert/blob/main/src/inference/batch.py): `src/inference/batch.py`
- [Checkpoint loader](https://github.com/lingyi0110/mai2bert/blob/main/src/inference/loader.py): `src/inference/loader.py`
- [Prediction runtime](https://github.com/lingyi0110/mai2bert/blob/main/src/inference/predict.py): `src/inference/predict.py`
- **Prediction files**

### Other components

- **Chart files**
- **Model checkpoint**

## Connections

- Chart files → Prepare command: supplies charts
- Prepare command → Preparation pipeline: runs preparation
- Preparation pipeline → Chart discovery: discovers charts
- Preparation pipeline → Simai parser: parses charts
- Preparation pipeline → Event representation: builds events
- Preparation pipeline → Chart augmentation: creates variants
- Preparation pipeline → Dataset splits: assigns splits
- Preparation pipeline → Processed HDF5: writes arrays
- Chart files → Inference command: supplies charts
- Inference command → Batch prediction: dispatches batch
- Batch prediction → Simai parser: parses Simai
- Batch prediction → MA2 parser: parses MA2
- Batch prediction → Prediction runtime: submits items
- Inference command → Checkpoint loader: loads model
- Model checkpoint → Checkpoint loader: provides weights
- Checkpoint loader → Difficulty regressor: restores model
- Prediction runtime → Difficulty regressor: predicts difficulty
- Batch prediction → Prediction files: writes CSV and JSON
- Training command → Pretraining loop: starts pretraining
- Training command → Regression training: starts regression
- Training command → Transfer training: starts transfer
- Processed HDF5 → Event dataset: supplies events
- Event dataset → Batch collation: provides samples
- Batch collation → Pretraining masking: collates for masking
- Pretraining masking → Pretraining model: provides masked batches
- Batch collation → Difficulty regressor: provides batches
- Pretraining model → Event embedding: encodes events
- Difficulty regressor → Event embedding: encodes events
- Event embedding → Sequence encoder: builds representations
- Pretraining loop → Pretraining model: trains model
- Regression training → Difficulty regressor: trains regressor
- Transfer training → Difficulty regressor: adapts model
- Regression training → Model checkpoint: saves weights