Update README.md

README.md

@@ -35,9 +35,10 @@ model_size = "5M"  # Options: "35M", "30M", "11M", "5M", "1.25M"
 model_config = MODEL_CONFIGS[model_size]
 
 # Load appropriate model
-model_path = f"chandan-sreedhara/SimpleStories-{model_size}"
+model_path = f"SimpleStories/SimpleStories-{model_size}"
 model = Llama.from_pretrained(model_path, model_config)
-model.to("cuda")
+device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
+model.to(device)
 model.eval()
 
 # Load tokenizer
@@ -47,14 +48,14 @@ tokenizer = AutoTokenizer.from_pretrained(model_path)
 prompt = "The curious cat looked at the"
 
 inputs = tokenizer(prompt, return_tensors="pt")
-input_ids = inputs.input_ids.to("cuda")
+input_ids = inputs.input_ids.to(device)
 
 # Generate text
 with torch.no_grad():
     output_ids = model.generate(
         idx=input_ids,
-        max_new_tokens=800,
-        temperature=0.7,
+        max_new_tokens=50,
+        temperature=0.0,
         top_k=40,
         eos_token_id=tokenizer.eos_token_id
     )
@@ -62,6 +63,7 @@ with torch.no_grad():
 # Decode output
 output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 print(f"Generated text:\n{output_text}")
+
 ```
 
 ## Model Variants