Update README.md

README.md

@@ -1,199 +1,205 @@
 ---
 library_name: transformers
-tags: []
+license: apache-2.0
+datasets:
+- ILSVRC/imagenet-1k
 ---
+# NEPA: Next-Embedding Predictive Architectures Are Strong Vision Learners
 
-# Model Card for Model ID
+[![Paper](https://img.shields.io/badge/arXiv-Paper-b31b1b?logo=arxiv&logoColor=b31b1b)]()
+[![Project Page](https://img.shields.io/badge/Project-Website-5B7493?logo=googlechrome&logoColor=5B7493)]()
+[![Hugging Model Card](https://img.shields.io/badge/huggingface-model:NEPA-lightblue)]()
 
-<!-- Provide a quick summary of what the model is/does. -->
+This is a PyTorch/GPU re-implementation of Next-Embedding Predictive Architectures Are Strong Vision Learners.
 
+<p align="center">
+  <img src="https://github.com/user-attachments/assets/760a7ca4-1aa4-43d9-8731-ab8b33e463d5" width="350">
+</p>
 
+The Next-Embedding Predictive Architecture. An image is split into patches and embedded into a sequence. An autoregressive model predicts the next embedding from previous ones.
 
-## Model Details
+```
+@article{six2025nepa,
+  title={Next-Embedding Predictive Architectures Are Strong Vision Learners},
+  author = {Sihan Xu and Ziqiao Ma and Wenhao Chai and Xuweiyi Chen and Weiyang Jin and Joyce Chai and Saining Xie and Stella X. Yu},
+  journal={arXiv preprint arXiv:},
+  year={2025}
+}
 
-### Model Description
+```
 
-<!-- Provide a longer summary of what this model is. -->
 
-This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+## Environment
 
-- **Developed by:** [More Information Needed]
-- **Funded by [optional]:** [More Information Needed]
-- **Shared by [optional]:** [More Information Needed]
-- **Model type:** [More Information Needed]
-- **Language(s) (NLP):** [More Information Needed]
-- **License:** [More Information Needed]
-- **Finetuned from model [optional]:** [More Information Needed]
+The codebase has been tested with the following environment:
 
-### Model Sources [optional]
+- Python 3.10
+- PyTorch 2.8.0
+- Transformers 4.56.2
 
-<!-- Provide the basic links for the model. -->
+### Installation
 
-- **Repository:** [More Information Needed]
-- **Paper [optional]:** [More Information Needed]
-- **Demo [optional]:** [More Information Needed]
+First, clone the repository:
 
-## Uses
+```bash
+git clone https://github.com/SihanXU/nepa
+cd nepa
+```
 
-<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+Then, create a conda environment and install dependencies:
 
-### Direct Use
+```bash
+conda env create -f environment.yml
+conda activate nepa
+```
 
-<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+Alternatively, you can install the dependencies manually:
 
-[More Information Needed]
+```bash
+pip install -r requirements.txt
+```
 
-### Downstream Use [optional]
+## Quick Start
 
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+Here's a simple example to run inference with a pretrained NEPA model:
 
-[More Information Needed]
+```python
+from transformers import AutoImageProcessor
+from models.vit_nepa import ViTNepaForImageClassification
+from PIL import Image
+import requests
 
-### Out-of-Scope Use
+url = 'https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg'
+image = Image.open(requests.get(url, stream=True).raw)
 
-<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+processor = AutoImageProcessor.from_pretrained('SixAILab/nepa-large-patch14-224-sft')
+model = ViTNepaForImageClassification.from_pretrained('SixAILab/nepa-large-patch14-224-sft')
 
-[More Information Needed]
+inputs = processor(images=image, return_tensors="pt")
+outputs = model(**inputs)
+logits = outputs.logits
+# model predicts one of the 1000 ImageNet classes
+predicted_class_idx = logits.argmax(-1).item()
+print("Predicted class:", model.config.id2label[predicted_class_idx])
+```
 
-## Bias, Risks, and Limitations
+## Setup Huggingface Token
 
-<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+To download pretrained models from Hugging Face Hub, you need to authenticate with your Hugging Face account:
 
-[More Information Needed]
+```bash
+hf auth login
+```
 
-### Recommendations
 
-<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+## [Optional] Setup Wandb Token
 
-Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+We use Wandb to track experiments. You may want to use [Weights & Biases](https://wandb.ai/) to log and track your experiments:
 
-## How to Get Started with the Model
+```bash
+pip install wandb
+wandb login
+```
 
-Use the code below to get started with the model.
+## Prepare ImageNet-1k Dataset
 
-[More Information Needed]
+We use the ImageNet-1k dataset for training and evaluation. To download the dataset via Hugging Face Datasets:
 
-## Training Details
+```bash
+python download_dataset.py
+```
 
-### Training Data
+This script will download and prepare the ImageNet-1k dataset. Note that this requires approximately 150GB of disk space. You may need to accept the dataset terms on [Hugging Face](https://huggingface.co/datasets/ILSVRC/imagenet-1k) before downloading.
 
-<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+## Evaluate Nepa for Image Classification
 
-[More Information Needed]
+We provide pretrained checkpoints for NEPA models. The following table compares our reproduced results with the paper:
 
-### Training Procedure
+| Model  | SwiGLU (paper) | GeLU (reproduce) |
+|--------|---------------:|-----------------:|
+| Nepa-B | 83.8           | 83.75            |
+| Nepa-L | 85.3           | 85.40            |
 
-<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+To evaluate the base model on ImageNet-1k validation set:
 
-#### Preprocessing [optional]
+```bash
+bash scripts/eval/nepa_b_sft_eval.sh 
+```
 
-[More Information Needed]
+This should give:
+```
+***** eval metrics *****
+  eval_accuracy               =     0.8375
+  eval_loss                   =     0.7169
+```
 
+To evaluate the large model:
 
-#### Training Hyperparameters
+```bash
+bash scripts/eval/nepa_l_sft_eval.sh 
+```
 
-- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+This should give:
+```
+***** eval metrics *****
+  eval_accuracy               =      0.854
+  eval_loss                   =     0.6371
+```
 
-#### Speeds, Sizes, Times [optional]
+## Fine-tune
 
-<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+To fine-tune a pretrained NEPA model on ImageNet-1k for image classification:
 
-[More Information Needed]
+For the base model:
 
-## Evaluation
+```bash
+bash scripts/finetune/nepa_b_sft.sh 
+```
 
-<!-- This section describes the evaluation protocols and provides the results. -->
+For the large model:
 
-### Testing Data, Factors & Metrics
+```bash
+bash scripts/finetune/nepa_l_sft.sh 
+```
 
-#### Testing Data
+You can modify the training hyperparameters (learning rate, batch size, epochs, etc.) in the corresponding script files.
 
-<!-- This should link to a Dataset Card if possible. -->
+## Pretrain
 
-[More Information Needed]
+To pretrain NEPA from scratch on ImageNet-1k:
 
-#### Factors
+For the base model:
 
-<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+```bash
+bash scripts/pretrain/nepa_b.sh
+```
 
-[More Information Needed]
+For the large model:
 
-#### Metrics
+```bash
+bash scripts/pretrain/nepa_l.sh
+```
 
-<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+Pretraining typically requires multiple GPUs. We recommend using at least 8 A100 GPUs for the large model.
 
-[More Information Needed]
+## Convert a Pretrained Model to Classification Model
 
-### Results
+After pretraining, you can convert the pretrained model to a classification model by initializing a classification head. Use the `init_nepa_cls_from_pretrain.py` script:
 
-[More Information Needed]
+Here is an example:
+```
+python init_nepa_cls_from_pretrain.py \
+  --pretrained_model_id SixAILab/nepa-base-patch14-224 \
+  --config_model_id configs/finetune/nepa-base-patch14-224-sft \
+  --pretrained_revision main \
+  --save_local \
+  --local_dir ./nepa-base-patch14-224-sft
+```
 
-#### Summary
+## Acknowledgements
 
+We gratefully acknowledge the developers of [Transformers](https://github.com/huggingface/transformers), [Evaluate](https://github.com/huggingface/evaluate), and [timm](https://github.com/huggingface/pytorch-image-models) for their excellent open-source contributions.
 
+## Contact
 
-## Model Examination [optional]
-
-<!-- Relevant interpretability work for the model goes here -->
-
-[More Information Needed]
-
-## Environmental Impact
-
-<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
-
-Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
-
-- **Hardware Type:** [More Information Needed]
-- **Hours used:** [More Information Needed]
-- **Cloud Provider:** [More Information Needed]
-- **Compute Region:** [More Information Needed]
-- **Carbon Emitted:** [More Information Needed]
-
-## Technical Specifications [optional]
-
-### Model Architecture and Objective
-
-[More Information Needed]
-
-### Compute Infrastructure
-
-[More Information Needed]
-
-#### Hardware
-
-[More Information Needed]
-
-#### Software
-
-[More Information Needed]
-
-## Citation [optional]
-
-<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
-
-**BibTeX:**
-
-[More Information Needed]
-
-**APA:**
-
-[More Information Needed]
-
-## Glossary [optional]
-
-<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
-
-[More Information Needed]
-
-## More Information [optional]
-
-[More Information Needed]
-
-## Model Card Authors [optional]
-
-[More Information Needed]
-
-## Model Card Contact
-
-[More Information Needed]
+Feel free to contact me through email (sihanxu@umich.edu). Enjoy!