Update README.md

README.md

@@ -21,7 +21,7 @@ library_name: transformers
 
 **NeoLM-Small-Base** is an English-focused, decoder-only Transformer language model trained from scratch with PyTorch and Hugging Face Transformers. This repository publishes the **Stage 1 Foundation** base-model checkpoint in **bfloat16** format. It is a **base next-token prediction model**, not an instruction-tuned chat assistant.
 
-The fixed Foundation release can be loaded using the `stage1-foundation` revision after the release tag is created.
+The fixed Foundation release can be loaded using the `stage1-foundation` revision.
 
 ## Architecture
 
@@ -82,7 +82,6 @@ This is a released training checkpoint rather than a claim of complete convergen
 Install compatible versions of `torch` and `transformers`. The model requires custom Python code, so review the repository code before enabling `trust_remote_code=True`.
 
 ```python
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 repo_id = "MaagDeveloper/NeoLM-Small-Base"
@@ -91,28 +90,79 @@ revision = "stage1-foundation"
 tokenizer = AutoTokenizer.from_pretrained(
     repo_id, revision=revision, trust_remote_code=True
 )
+
 model = AutoModelForCausalLM.from_pretrained(
     repo_id, revision=revision, trust_remote_code=True,
-    dtype=torch.bfloat16,
+    dtype="auto",
 ).eval()
 
-# For a CPU that does not support BF16 well, load with dtype=torch.float32 instead.
+# For a CPU that does not support BF16 well, use dtype=torch.float32 instead.
 prompt = "The future of artificial intelligence"
-inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
-with torch.no_grad():
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=64,
-        do_sample=True,
-        temperature=0.8,
-        top_p=0.9,
-        repetition_penalty=1.15,
-        eos_token_id=tokenizer.eos_token_id,
-        pad_token_id=tokenizer.pad_token_id,
-    )
+
+inputs = tokenizer(
+    prompt, return_tensors="pt", truncation=True, max_length=1024
+)
+
+outputs = model.generate(
+    **inputs,
+    max_new_tokens=64,
+    do_sample=True,
+    temperature=0.8,
+    top_p=0.9,
+    repetition_penalty=1.15,
+    eos_token_id=tokenizer.eos_token_id,
+    pad_token_id=tokenizer.pad_token_id,
+)
+
 print(tokenizer.decode(outputs[0], skip_special_tokens=True))
 ```
 
+
+## Streaming Text Generation
+
+You can stream generated text progressively using Hugging Face Transformers' `TextStreamer`.
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+
+repo_id = "MaagDeveloper/NeoLM-Small-Base"
+revision = "stage1-foundation"
+
+tokenizer = AutoTokenizer.from_pretrained(
+    repo_id,
+    revision=revision,
+    trust_remote_code=True,
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    repo_id,
+    revision=revision,
+    trust_remote_code=True,
+    dtype="auto",
+).eval()
+
+streamer = TextStreamer(
+    tokenizer,
+    skip_prompt=False,
+    skip_special_tokens=True,
+)
+
+prompt = "Once upon a time"
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+model.generate(
+    **inputs,
+    streamer=streamer,
+    max_new_tokens=100,
+    do_sample=True,
+    temperature=0.8,
+    top_p=0.9,
+    repetition_penalty=1.15,
+    eos_token_id=tokenizer.eos_token_id,
+    pad_token_id=tokenizer.pad_token_id,
+)
+```
+
 The code above is **text completion**, not role-based chat. NeoLM-Small-Base may not reliably follow instructions, answer factual questions, or perform multi-step reasoning. A separate instruction-tuned model belongs in a distinct repository.
 
 ## Evaluation and limitations