Update README.md

README.md

@@ -1,18 +1,16 @@
 ---
-license: mit
+license: cc-by-nc-4.0
 ---
 
-# Model card for boldgpt_small_patch10
+# Model card for `boldgpt_small_patch10.kmq`
 
 ![Example training predictions](example.png)
 
-A Vision Transformer (ViT) model trained on BOLD activation maps from [NSD-Flat](https://huggingface.co/datasets/clane9/NSD-Flat). The training objective was to auto-regressively predict the next patch with shuffled patch order.
+A Vision Transformer (ViT) model trained on BOLD activation maps from [NSD-Flat](https://huggingface.co/datasets/clane9/NSD-Flat). Patches were quantized to discrete tokens using k-means (`KMeansTokenizer`). The training objective was to auto-regressively predict the next patch with shuffled patch order and cross-entropy loss.
 
 ## Dependencies
 
 - [boldGPT](https://github.com/clane9/boldGPT)
-- [huggingface_hub](https://huggingface.co/docs/huggingface_hub/index)
-- [safetensors](https://huggingface.co/docs/safetensors/index)
 
 ## Usage
 
@@ -20,25 +18,16 @@ A Vision Transformer (ViT) model trained on BOLD activation maps from [NSD-Flat]
 from boldgpt.data import ActivityTransform
 from boldgpt.models import create_model
 from datasets import load_dataset
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_model
 
-model = create_model("boldgpt_small_patch10")
-
-load_model(
-    model,
-    hf_hub_download(
-        repo_id="clane9/boldgpt_small_patch10", filename="model.safetensors"
-    ),
-)
+model = create_model("boldgpt_small_patch10.kmq", pretrained=True)
 
 dataset = load_dataset("clane9/NSD-Flat", split="train")
 dataset.set_format("torch")
-batch = dataset[:1]
 
 transform = ActivityTransform()
+batch = dataset[:1]
 batch["activity"] = transform(batch["activity"])
 
-# output: (B, N, K) predicted next token logits
+# output: (B, N + 1, K) predicted next token logits
 output, state = model(batch)
 ```