Update README.md

README.md

@@ -16,33 +16,23 @@ Interpretable slot-based attention achieving competitive language modeling perfo
 ## Quick Start
 
 ```python
-!wget https://huggingface.co/DigitalDaimyo/AddressedStateAttention/resolve/main/asm_analysis.py
-!wget https://huggingface.co/DigitalDaimyo/AddressedStateAttention/resolve/main/universal_loader.py
+# Install directly from GitHub
+!pip install git+https://github.com/DigitalDaimyo/AddressedStateAttention.git
 
-from universal_loader import load_asm_checkpoint
-from huggingface_hub import hf_hub_download
+from asa import load_asm_checkpoint, generate
+from transformers import AutoTokenizer
 
-# Download checkpoint
-ckpt_path = hf_hub_download(
-    repo_id="DigitalDaimyo/AddressedStateAttention",
-    filename="checkpoints/fineweb_187M_75k.pt"
+# Load checkpoint
+model, cfg, ckpt = load_asm_checkpoint(
+    "path/to/checkpoint.pt",
+    mode="analysis"
 )
 
-# Load model
-model, cfg, ckpt = load_asm_checkpoint(ckpt_path, mode="analysis")
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
 
 # Generate text
-from transformers import AutoTokenizer
-from generation import generate
+print(generate(model, tokenizer, "Once upon a time"))
 
-tokenizer = AutoTokenizer.from_pretrained("gpt2")
-text = generate(
-    model, tokenizer,
-    prompt="The capital of France is",
-    max_new_tokens=20,
-    strategy="greedy"
-)
-print(text)
 Performance
 FineWeb, 187M params: 3.73 val loss / 41.6 PPL (75k steps)
 Architecture: 21 layers, 768d, 12 heads, 16 slots