How meetmendapara09/attenta works
Purpose and entry points
Attenta is a standalone Rust crate implementing and training a Transformer language model from scratch on Indian legal-act JSON corpora. The executable entry point is transformer/src/main.rs; reusable crate exports and module wiring live in transformer/src/lib.rs. A dedicated legal-training example is transformer/examples/train_legal.rs.
Corpus and input pipeline
Legal source material is stored as JSON files under data/, including data/The Code on Wages 2019.json and data/The Insolvency And Bankruptcy Code 2016.json. Data ingestion and dataset preparation are centered in transformer/src/data.rs, while BPE-style token vocabulary/encoding is in transformer/src/tokenizer.rs. transformer/src/batch.rs groups tokenized examples into length-aware batches, forming the boundary between corpus processing and numerical training.
Transformer model core
transformer/src/model.rs composes the encoder/decoder Transformer and model-level forward behavior. Encoder stacks are implemented in transformer/src/encoder.rs; autoregressive decoder stacks and target-side masking behavior are in transformer/src/decoder.rs. These modules represent the primary sequence-to-sequence execution path over token IDs and embeddings.
Attention and tensor execution
Multi-head scaled dot-product attention is implemented in transformer/src/attention.rs. Core tensor/matrix operations, including performance-oriented matrix multiplication, reside in transformer/src/tensor_ops.rs. This is the main numerical boundary: higher-level model modules delegate array math and attention score transformations to these lower-level routines.
Training, gradients, and optimization
transformer/src/train.rs orchestrates the training loop: batches flow through model forward passes, loss computation, backpropagation, and parameter updates. Manual gradient propagation is implemented in transformer/src/backward.rs. Objective calculation, including label-smoothed cross entropy, is in transformer/src/loss.rs; Adam and learning-rate warmup scheduling are in transformer/src/optim.rs.
Decoding and model artifacts
Inference generation supports greedy and beam-search decoding through transformer/src/decoder.rs and associated model interfaces in transformer/src/model.rs. Training state and model parameter persistence are handled by transformer/src/checkpoint.rs, with committed artifacts at transformer/checkpoint.json and transformer/checkpoint_avg.json. Attention weights can be emitted or inspected through transformer/src/visualize.rs and transformer/attention_weights.json.
Runtime and boundaries
The system runs as a native Rust 1.70+ process via Cargo, with no PyTorch, TensorFlow, hosted service, database, or web runtime indicated. Execution may require an increased Rust stack (RUST_MIN_STACK) for the release binary. The only material external input is the local legal-document corpus; checkpoints and visualizations are local JSON artifacts.
Main entry points: model.rs Transformer model · train.rs Training loop · decoder.rs Decoder and generation · attention.rs Multi-head attention
Components 20
Entry points
- CLI executable · Native Cargo entry point.
transformer/src/main.rs - Crate exports · Module wiring and reusable interfaces.
transformer/src/lib.rs - Legal training example
transformer/examples/train_legal.rs
Corpus pipeline
- Legal-act JSON corpus · Indian legal documents under the data directory.
- Dataset preparation
transformer/src/data.rs - BPE tokenizer
transformer/src/tokenizer.rs - Length-aware batching · Boundary from tokenized examples to numerical training batches.
transformer/src/batch.rs
Transformer core
- Transformer model · Composes encoder, decoder, and model forward behavior.
transformer/src/model.rs - Encoder stacks
transformer/src/encoder.rs - Decoder and generation · Autoregressive stacks, masking, greedy and beam decoding.
transformer/src/decoder.rs - Multi-head attention · Scaled dot-product attention.
transformer/src/attention.rs - Tensor operations · Matrix operations and performance-oriented multiplication.
transformer/src/tensor_ops.rs
Training runtime
- Training loop · Coordinates forward pass, loss, backpropagation, and updates.
transformer/src/train.rs - Label-smoothed loss
transformer/src/loss.rs - Manual backpropagation
transformer/src/backward.rs - Adam with warmup
transformer/src/optim.rs
Local artifacts
- Checkpoint persistence · Saves and restores training state and parameters.
transformer/src/checkpoint.rs - Model checkpoints
transformer/checkpoint.json - Attention weights · Inspectable emitted attention weights.
transformer/attention_weights.json - Attention visualization
transformer/src/visualize.rs
Connections 21
- CLI executable → Crate exports: uses
- Legal training example → Training loop: runs
- Legal-act JSON corpus → Dataset preparation: ingests
- Dataset preparation → BPE tokenizer: prepares text for
- BPE tokenizer → Length-aware batching: token IDs
- Length-aware batching → Training loop: batches
- Training loop → Transformer model: forward pass
- Transformer model → Encoder stacks: encodes source
- Transformer model → Decoder and generation: decodes targets
- Encoder stacks → Multi-head attention: uses
- Decoder and generation → Multi-head attention: uses masked attention
- Multi-head attention → Tensor operations: score and matrix math
- Training loop → Label-smoothed loss: computes objective
- Label-smoothed loss → Manual backpropagation: loss gradients
- Manual backpropagation → Adam with warmup: parameter gradients
- Adam with warmup → Transformer model: updates parameters
- Training loop → Checkpoint persistence: saves/restores state
- Checkpoint persistence → Model checkpoints: writes
- Transformer model → Attention visualization: attention data
- Attention visualization → Attention weights: emits
- Decoder and generation → Transformer model: generation interface