# ganeshedula/capstone architecture

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

- Interactive diagram: https://gitdiagram.com/ganeshedula/capstone
- Repository: https://github.com/ganeshedula/capstone

## Overview

This repository is a DoLa + ENN hallucination-reduction and evaluation pipeline for TinyLlama. Its principal workflow is CLI invocation, model loading, C4 feature preparation, ENN training, then model evaluation and comparison plots; it also offers inference and DoLa experiments. The source index identifies substantial modules and some references, but the README is empty and excerpts omit key orchestration and subsystem internals. Accordingly, uncertain main.py dispatches are shown as unverified rather than asserted.

Main entry points: [`main.py`](https://github.com/ganeshedula/capstone/blob/main/main.py) (CLI pipeline), [`Neural Network/enn_torch.py`](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/enn_torch.py) (Torch Epinet), [`Neural Network/before_after_eval.py`](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/before_after_eval.py) (Before/after evaluation), [`Neural Network/c4_eval.py`](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/c4_eval.py) (C4 token evaluation)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_orchestration["Pipeline entry"]
  node_cli["CLI pipeline<br/>[main.py]"]
end

subgraph group_modeling["Model and decoding"]
  node_model_loader["TinyLlama loader<br/>[load_llama2.py]"]
  node_dola["DoLa contrastive decoding<br/>[dola.py]"]
  node_inference["Autoregressive inference<br/>[inference.py]"]
end

subgraph group_training["ENN preparation and training"]
  node_c4_data["C4 feature preparation<br/>[data_prep.py]"]
  node_enn_torch["Torch Epinet<br/>[enn_torch.py]"]
  node_enn_jax["JAX ENN<br/>[enn_jax.py]"]
end

subgraph group_evaluation["Evaluation and analysis"]
  node_truthful_eval["TruthfulQA evaluation<br/>[evaluate.py]"]
  node_before_after["Before/after evaluation"]
  node_before_plots["Before/after plots"]
  node_c4_eval["C4 token evaluation<br/>[c4_eval.py]"]
  node_c4_plots["C4 evaluation plots<br/>[c4_eval_plots.py]"]
end

node_c4_source[("C4 text source")]
node_operator(("CLI user"))

node_operator -->|"invokes"| node_cli
node_cli -.->|"loads model"| node_model_loader
node_cli -.->|"prepares features"| node_c4_data
node_cli -.->|"trains ENN"| node_enn_torch
node_cli -.->|"runs evaluation"| node_truthful_eval
node_cli -.->|"runs comparison"| node_before_after
node_cli -.->|"generates plots"| node_before_plots
node_cli -.->|"runs C4 evaluation"| node_c4_eval
node_cli -.->|"generates plots"| node_c4_plots
node_c4_source -->|"provides text"| node_c4_data
node_c4_source -.->|"provides held-out text"| node_c4_eval
node_c4_data -->|"extracts hidden states"| node_dola
node_c4_data -.->|"produces training features"| node_enn_torch
node_truthful_eval -->|"scores with DoLa"| node_dola
node_truthful_eval -->|"uses ENN logits"| node_enn_torch
node_before_after -->|"reuses scoring"| node_truthful_eval
node_before_after -->|"evaluates DoLa"| node_dola
node_before_after -->|"evaluates DoLa plus ENN"| node_enn_torch
node_before_plots -->|"plots results"| node_before_after
node_c4_eval -->|"scores tokens"| node_dola
node_c4_eval -->|"scores with ENN"| node_enn_torch
node_c4_plots -->|"plots results"| node_c4_eval
node_inference -->|"gets contrastive logits"| node_dola
node_inference -->|"gets ENN predictions"| node_enn_torch

click node_cli "https://github.com/ganeshedula/capstone/blob/main/main.py"
click node_model_loader "https://github.com/ganeshedula/capstone/blob/main/Model_load/load_llama2.py"
click node_c4_data "https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/data_prep.py"
click node_dola "https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/dola.py"
click node_enn_torch "https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/enn_torch.py"
click node_enn_jax "https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/enn_jax.py"
click node_inference "https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/inference.py"
click node_truthful_eval "https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/evaluate.py"
click node_before_after "https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/before_after_eval.py"
click node_before_plots "https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/before_after_plots.py"
click node_c4_eval "https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/c4_eval.py"
click node_c4_plots "https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/c4_eval_plots.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_cli,node_operator toneBlue
class node_model_loader,node_dola,node_inference,node_c4_source toneAmber
class node_c4_data,node_enn_torch,node_enn_jax toneMint
class node_truthful_eval,node_before_after,node_before_plots,node_c4_eval,node_c4_plots toneRose
```

## Components

### Pipeline entry

- [CLI pipeline](https://github.com/ganeshedula/capstone/blob/main/main.py): `main.py`

### Model and decoding

- [TinyLlama loader](https://github.com/ganeshedula/capstone/blob/main/Model_load/load_llama2.py): `Model_load/load_llama2.py`
- [DoLa contrastive decoding](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/dola.py): `Neural Network/dola.py`
- [Autoregressive inference](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/inference.py): `Neural Network/inference.py`

### ENN preparation and training

- [C4 feature preparation](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/data_prep.py): `Neural Network/data_prep.py`
- [Torch Epinet](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/enn_torch.py): `Neural Network/enn_torch.py`
- [JAX ENN](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/enn_jax.py): `Neural Network/enn_jax.py`

### Evaluation and analysis

- [TruthfulQA evaluation](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/evaluate.py): `Neural Network/evaluate.py`
- [Before/after evaluation](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/before_after_eval.py): `Neural Network/before_after_eval.py`
- [Before/after plots](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/before_after_plots.py): `Neural Network/before_after_plots.py`
- [C4 token evaluation](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/c4_eval.py): `Neural Network/c4_eval.py`
- [C4 evaluation plots](https://github.com/ganeshedula/capstone/blob/main/Neural%20Network/c4_eval_plots.py): `Neural Network/c4_eval_plots.py`

### Other components

- **C4 text source**
- **CLI user**

## Connections

- CLI user → CLI pipeline: invokes (evidence: `main.py`)
- CLI pipeline → TinyLlama loader: loads model (evidence: `main.py`)
- CLI pipeline → C4 feature preparation: prepares features (evidence: `main.py`)
- CLI pipeline → Torch Epinet: trains ENN (evidence: `main.py`)
- CLI pipeline → TruthfulQA evaluation: runs evaluation (evidence: `main.py`)
- CLI pipeline → Before/after evaluation: runs comparison (evidence: `main.py`)
- CLI pipeline → Before/after plots: generates plots (evidence: `main.py`)
- CLI pipeline → C4 token evaluation: runs C4 evaluation (evidence: `main.py`)
- CLI pipeline → C4 evaluation plots: generates plots (evidence: `main.py`)
- C4 text source → C4 feature preparation: provides text (evidence: `Neural Network/data_prep.py`)
- C4 text source → C4 token evaluation: provides held-out text
- C4 feature preparation → DoLa contrastive decoding: extracts hidden states (evidence: `Neural Network/data_prep.py`)
- C4 feature preparation → Torch Epinet: produces training features
- TruthfulQA evaluation → DoLa contrastive decoding: scores with DoLa (evidence: `Neural Network/evaluate.py`)
- TruthfulQA evaluation → Torch Epinet: uses ENN logits (evidence: `Neural Network/evaluate.py`)
- Before/after evaluation → TruthfulQA evaluation: reuses scoring (evidence: `Neural Network/before_after_eval.py`)
- Before/after evaluation → DoLa contrastive decoding: evaluates DoLa (evidence: `Neural Network/before_after_eval.py`)
- Before/after evaluation → Torch Epinet: evaluates DoLa plus ENN (evidence: `Neural Network/before_after_eval.py`)
- Before/after plots → Before/after evaluation: plots results (evidence: `Neural Network/before_after_plots.py`)
- C4 token evaluation → DoLa contrastive decoding: scores tokens (evidence: `Neural Network/c4_eval.py`)
- C4 token evaluation → Torch Epinet: scores with ENN (evidence: `Neural Network/c4_eval.py`)
- C4 evaluation plots → C4 token evaluation: plots results (evidence: `Neural Network/c4_eval_plots.py`)
- Autoregressive inference → DoLa contrastive decoding: gets contrastive logits (evidence: `Neural Network/inference.py`)
- Autoregressive inference → Torch Epinet: gets ENN predictions (evidence: `Neural Network/inference.py`)