# nithish-007/mi_jepa_based_clustering architecture

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

- Interactive diagram: https://gitdiagram.com/nithish-007/mi_jepa_based_clustering
- Repository: https://github.com/nithish-007/mi_jepa_based_clustering

## Overview

This repository implements an SKU/product-image clustering workflow: prepare image data, train a representation model, cluster unseen products, and evaluate clustering and retrieval quality. The model combines a ViT backbone, invariance and discrimination heads, and an optional I-JEPA training branch; inference embeds images and offers graph, HDBSCAN, and oracle k-means methods. Data-module implementations and some pipeline details are not shown, so their internal connections are omitted; edges reflect the supplied excerpts, source index, or file names.

Main entry points: [`src/models/network.py`](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/models/network.py) (Clustering network), [`src/training/module.py`](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/training/module.py) (Training module), [`src/evaluation/pipeline.py`](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/evaluation/pipeline.py) (Evaluation pipeline), [`src/inference/embedding.py`](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/embedding.py) (Model embedding)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_workflow["Product workflow"]
  node_config["Config loader<br/>[config.py]"]
  node_download["RP2K download<br/>[download_rp2k.py]"]
  node_prepare["Data preparation<br/>[prepare_data.py]"]
  node_data["Image data"]
  node_train_cli["Train command<br/>[train.py]"]
  node_cluster_cli["Cluster command<br/>[cluster.py]"]
  node_eval_cli["Evaluate command<br/>[evaluate.py]"]
end

subgraph group_model["Model and training"]
  node_trainer["Training runner<br/>[trainer.py]"]
  node_module["Training module<br/>[module.py]"]
  node_backbone["ViT backbone<br/>[backbone.py]"]
  node_network["Clustering network<br/>[network.py]"]
  node_heads["Embedding heads<br/>[heads.py]"]
  node_ijepa["I-JEPA branch<br/>[ijepa.py]"]
  node_losses["Training losses"]
  node_checkpoints["Checkpoint handling<br/>[checkpoints.py]"]
end

subgraph group_inference["Inference and clustering"]
  node_run["Inference runner<br/>[run.py]"]
  node_embed["Model embedding<br/>[embedding.py]"]
  node_inference_pipeline["Method dispatch<br/>[pipeline.py]"]
  node_scoring["Pair scoring<br/>[scoring.py]"]
  node_clustering["Clustering methods<br/>[clustering.py]"]
  node_export["Cluster export<br/>[export.py]"]
end

subgraph group_evaluation["Evaluation"]
  node_evaluation_pipeline["Evaluation pipeline<br/>[pipeline.py]"]
  node_metrics["Metrics<br/>[metrics.py]"]
  node_plots["Evaluation plots<br/>[plots.py]"]
  node_analysis["Model analysis<br/>[analysis.py]"]
end

node_user(("Researcher"))

node_user -->|"obtains data"| node_download
node_user -->|"prepares data"| node_prepare
node_user -->|"starts training"| node_train_cli
node_user -->|"runs clustering"| node_cluster_cli
node_user -->|"runs evaluation"| node_eval_cli
node_prepare -->|"loads settings"| node_config
node_prepare -.->|"prepares images"| node_data
node_download -.->|"downloads images"| node_data
node_train_cli -->|"loads settings"| node_config
node_train_cli -->|"starts training"| node_trainer
node_trainer -->|"trains module"| node_module
node_module -->|"forwards images"| node_network
node_module -->|"computes objectives"| node_losses
node_network -->|"extracts features"| node_backbone
node_network -->|"projects features"| node_heads
node_network -->|"creates branch"| node_ijepa
node_trainer -->|"manages checkpoints"| node_checkpoints
node_cluster_cli -->|"loads settings"| node_config
node_cluster_cli -->|"starts inference"| node_run
node_run -->|"loads model"| node_embed
node_embed -->|"builds baseline"| node_backbone
node_embed -->|"builds trained model"| node_network
node_embed -->|"resolves checkpoint"| node_checkpoints
node_run -->|"dispatches methods"| node_inference_pipeline
node_inference_pipeline -->|"runs baselines"| node_clustering
node_inference_pipeline -->|"scores graph pairs"| node_scoring
node_scoring -->|"clusters candidates"| node_clustering
node_run -->|"exports results"| node_export
node_eval_cli -->|"loads settings"| node_config
node_eval_cli -->|"starts evaluation"| node_evaluation_pipeline
node_evaluation_pipeline -->|"embeds products"| node_embed
node_evaluation_pipeline -->|"predicts clusters"| node_inference_pipeline
node_evaluation_pipeline -->|"scores predictions"| node_metrics
node_evaluation_pipeline -->|"saves diagnostics"| node_plots
node_module -->|"validates retrieval"| node_metrics
node_module -->|"fits calibration"| node_scoring
node_analysis -->|"uses metrics"| node_metrics
node_analysis -->|"uses clustering"| node_clustering
node_analysis -->|"analyzes model"| node_network
node_analysis -->|"analyzes objectives"| node_losses

click node_config "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/config.py"
click node_download "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/scripts/download_rp2k.py"
click node_prepare "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/scripts/prepare_data.py"
click node_data "https://github.com/nithish-007/mi_jepa_based_clustering/tree/main/src/data"
click node_train_cli "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/scripts/train.py"
click node_trainer "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/training/trainer.py"
click node_module "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/training/module.py"
click node_backbone "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/models/backbone.py"
click node_network "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/models/network.py"
click node_heads "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/models/heads.py"
click node_ijepa "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/models/ijepa.py"
click node_losses "https://github.com/nithish-007/mi_jepa_based_clustering/tree/main/src/losses"
click node_checkpoints "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/training/checkpoints.py"
click node_cluster_cli "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/scripts/cluster.py"
click node_run "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/run.py"
click node_embed "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/embedding.py"
click node_inference_pipeline "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/pipeline.py"
click node_scoring "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/scoring.py"
click node_clustering "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/clustering.py"
click node_export "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/export.py"
click node_eval_cli "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/scripts/evaluate.py"
click node_evaluation_pipeline "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/evaluation/pipeline.py"
click node_metrics "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/evaluation/metrics.py"
click node_plots "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/evaluation/plots.py"
click node_analysis "https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/evaluation/analysis.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_config,node_download,node_prepare,node_data,node_train_cli,node_cluster_cli,node_eval_cli toneBlue
class node_trainer,node_module,node_backbone,node_network,node_heads,node_ijepa,node_losses,node_checkpoints toneAmber
class node_run,node_embed,node_inference_pipeline,node_scoring,node_clustering,node_export toneMint
class node_evaluation_pipeline,node_metrics,node_plots,node_analysis toneRose
class node_user toneIndigo
```

## Components

### Product workflow

- [Config loader](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/config.py): `src/config.py`
- [RP2K download](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/scripts/download_rp2k.py): `scripts/download_rp2k.py`
- [Data preparation](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/scripts/prepare_data.py): `scripts/prepare_data.py`
- [Image data](https://github.com/nithish-007/mi_jepa_based_clustering/tree/main/src/data): `src/data`
- [Train command](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/scripts/train.py): `scripts/train.py`
- [Cluster command](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/scripts/cluster.py): `scripts/cluster.py`
- [Evaluate command](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/scripts/evaluate.py): `scripts/evaluate.py`

### Model and training

- [Training runner](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/training/trainer.py): `src/training/trainer.py`
- [Training module](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/training/module.py): `src/training/module.py`
- [ViT backbone](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/models/backbone.py): `src/models/backbone.py`
- [Clustering network](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/models/network.py): `src/models/network.py`
- [Embedding heads](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/models/heads.py): `src/models/heads.py`
- [I-JEPA branch](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/models/ijepa.py): `src/models/ijepa.py`
- [Training losses](https://github.com/nithish-007/mi_jepa_based_clustering/tree/main/src/losses): `src/losses`
- [Checkpoint handling](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/training/checkpoints.py): `src/training/checkpoints.py`

### Inference and clustering

- [Inference runner](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/run.py): `src/inference/run.py`
- [Model embedding](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/embedding.py): `src/inference/embedding.py`
- [Method dispatch](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/pipeline.py): `src/inference/pipeline.py`
- [Pair scoring](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/scoring.py): `src/inference/scoring.py`
- [Clustering methods](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/clustering.py): `src/inference/clustering.py`
- [Cluster export](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/inference/export.py): `src/inference/export.py`

### Evaluation

- [Evaluation pipeline](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/evaluation/pipeline.py): `src/evaluation/pipeline.py`
- [Metrics](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/evaluation/metrics.py): `src/evaluation/metrics.py`
- [Evaluation plots](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/evaluation/plots.py): `src/evaluation/plots.py`
- [Model analysis](https://github.com/nithish-007/mi_jepa_based_clustering/blob/main/src/evaluation/analysis.py): `src/evaluation/analysis.py`

### Other components

- **Researcher**

## Connections

- Researcher → RP2K download: obtains data
- Researcher → Data preparation: prepares data
- Researcher → Train command: starts training
- Researcher → Cluster command: runs clustering
- Researcher → Evaluate command: runs evaluation
- Data preparation → Config loader: loads settings (evidence: `scripts/prepare_data.py`)
- Data preparation → Image data: prepares images
- RP2K download → Image data: downloads images
- Train command → Config loader: loads settings (evidence: `scripts/train.py`)
- Train command → Training runner: starts training (evidence: `scripts/train.py`)
- Training runner → Training module: trains module (evidence: `src/training/trainer.py`)
- Training module → Clustering network: forwards images (evidence: `src/training/module.py`)
- Training module → Training losses: computes objectives (evidence: `src/training/module.py`)
- Clustering network → ViT backbone: extracts features (evidence: `src/models/network.py`)
- Clustering network → Embedding heads: projects features (evidence: `src/models/network.py`)
- Clustering network → I-JEPA branch: creates branch (evidence: `src/models/network.py`)
- Training runner → Checkpoint handling: manages checkpoints (evidence: `src/training/trainer.py`)
- Cluster command → Config loader: loads settings (evidence: `scripts/cluster.py`)
- Cluster command → Inference runner: starts inference (evidence: `scripts/cluster.py`)
- Inference runner → Model embedding: loads model (evidence: `src/inference/run.py`)
- Model embedding → ViT backbone: builds baseline (evidence: `src/inference/embedding.py`)
- Model embedding → Clustering network: builds trained model (evidence: `src/inference/embedding.py`)
- Model embedding → Checkpoint handling: resolves checkpoint (evidence: `src/inference/embedding.py`)
- Inference runner → Method dispatch: dispatches methods (evidence: `src/inference/run.py`)
- Method dispatch → Clustering methods: runs baselines (evidence: `src/inference/pipeline.py`)
- Method dispatch → Pair scoring: scores graph pairs (evidence: `src/inference/pipeline.py`)
- Pair scoring → Clustering methods: clusters candidates (evidence: `src/inference/scoring.py`)
- Inference runner → Cluster export: exports results (evidence: `src/inference/run.py`)
- Evaluate command → Config loader: loads settings (evidence: `scripts/evaluate.py`)
- Evaluate command → Evaluation pipeline: starts evaluation (evidence: `scripts/evaluate.py`)
- Evaluation pipeline → Model embedding: embeds products (evidence: `src/evaluation/pipeline.py`)
- Evaluation pipeline → Method dispatch: predicts clusters (evidence: `src/evaluation/pipeline.py`)
- Evaluation pipeline → Metrics: scores predictions (evidence: `src/evaluation/pipeline.py`)
- Evaluation pipeline → Evaluation plots: saves diagnostics (evidence: `src/evaluation/pipeline.py`)
- Training module → Metrics: validates retrieval (evidence: `src/training/module.py`)
- Training module → Pair scoring: fits calibration (evidence: `src/training/module.py`)
- Model analysis → Metrics: uses metrics (evidence: `src/evaluation/analysis.py`)
- Model analysis → Clustering methods: uses clustering (evidence: `src/evaluation/analysis.py`)
- Model analysis → Clustering network: analyzes model (evidence: `src/evaluation/analysis.py`)
- Model analysis → Training losses: analyzes objectives (evidence: `src/evaluation/analysis.py`)