Update README.md

README.md

@@ -17,52 +17,39 @@ The SimpleStories models are a tiny model family created for interpretability re
 
 ## Usage
 
-```bash
-pip install simple_stories_train
-```
-
 ```python
-from transformers import AutoTokenizer
 import torch
+from transformers import AutoTokenizer, LlamaForCausalLM
 
-from simple_stories_train.models.llama import Llama
-from simple_stories_train.models.model_configs import MODEL_CONFIGS
-
-# Select the model size you want to use
-model_size = "11M"  # Options: "35M", "30M", "11M", "5M", "1.25M"
 
-# Load model configuration
-model_config = MODEL_CONFIGS[model_size]
+MODEL_SIZE = "11M"
+model_path = "SimpleStories/SimpleStories-{}".format(MODEL_SIZE)
 
-# Load appropriate model
-model_path = f"SimpleStories/SimpleStories-{model_size}"
-model = Llama.from_pretrained(model_path, model_config)
-device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
-model.to(device)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = LlamaForCausalLM.from_pretrained(model_path)
+model.to("cuda")
 model.eval()
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_path)
 
-# Define your prompt
 prompt = "The curious cat looked at the"
 
-inputs = tokenizer(prompt, return_tensors="pt")
-input_ids = inputs.input_ids.to(device)
 
-# Generate text
+inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
+input_ids = inputs.input_ids.to("cuda")
+
+eos_token_id = 1
+
 with torch.no_grad():
     output_ids = model.generate(
-        idx=input_ids,
-        max_new_tokens=50,
-        temperature=0.0,
-        top_k=40,
-        eos_token_id=tokenizer.eos_token_id
-    )
-
-# Decode output
+        input_ids=input_ids,
+        max_new_tokens=400,
+        temperature=0.7,
+        do_sample=True,
+        eos_token_id=eos_token_id
+)
+
 output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-print(f"Generated text:\n{output_text}")
+print(f"\nGenerated text:\n{output_text}")
 
 ```
 
@@ -99,3 +86,7 @@ The SimpleStories dataset is a collection of short stories generated by state-of
 - ASCII-only guarantee for the English dataset
 
 Read the dataset paper on [arXiv](https://arxiv.org/abs/2504.09184).
+
+## Training
+
+The training and evaluation scripts can be accessed at https://github.com/danbraunai/simple_stories_train