Update README.md

README.md

@@ -2,82 +2,191 @@
 language: en
 license: apache-2.0
 library_name: transformers
+pipeline_tag: text-generation
 tags:
   - fwkv
   - rosa
+  - rwkv
   - linear-transformer
   - recurrent-neural-network
   - chat
   - ultrachat
-  - pytorch
+  - custom-architecture
 ---
 
 # FWKV-ROSA
 
-A **56M‑parameter**, from‑scratch recurrent language model combining a **per‑channel leaky integrator (FWKV)** with the **RWKV‑8 ROSA copy‑signal** mechanism.  
-Chat‑tuned on `HuggingFaceH4/ultrachat_200k` with a simple two‑role template.
+A **56M‑parameter**, from‑scratch recurrent language model that combines a
+**per‑channel leaky integrator (FWKV)** with the **RWKV‑8 ROSA copy‑signal**
+mechanism. It is a research experiment designed to explore how far purely
+recurrent architectures can go on small‑scale, curated conversational data.
 
-## Model description
+The model was chat‑tuned on `HuggingFaceH4/ultrachat_200k` and uses a simple
+two‑role template:
 
-- **Architecture:** 14‑layer FWKV blocks with a factorised tied embedding/head (512‑dim model, 128‑dim embedding) and an exact vectorised parallel scan for the constant‑decay recurrence.
-- **ROSA (Rapid Online Suffix Automaton):** A parameter‑free, causal predictor that injects the token that historically followed the longest matching suffix of the current context. Implemented as an embedding added to the early hidden state.
-- **Training data:** [UltraChat 200k](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k) (`train_sft` / `test_sft` splits), 20k training examples, 1k validation examples.
-- **Context length:** 1024 tokens (no positional embeddings – context is free parameter‑wise).
-- **Tokenizer:** GPT‑2 tokenizer extended with `<|user|>` and `<|assistant|>` special tokens.
-- **Training precision:** bfloat16 mixed precision on a single NVIDIA T4 (batch size 8, gradient accumulation 4, effective batch 32). Training took ~2 hours.
+```
+<|user|> Your message
+<|assistant|> Model reply<|endoftext|>
+```
+
+## Model Description
+
+- **Architecture:** 14 stacked FWKV blocks. Each block replaces the standard attention with a fixed, data‑independent decayed accumulator:
+  `stateₜ = W·stateₜ₋₁ + kₜ·vₜ`, where `W = clamp(sigmoid(w), min=0.1)`.
+  The recurrence is computed exactly via a vectorised parallel scan (no
+  approximations).
+- **ROSA (Rapid Online Suffix Automaton):** A parameter‑free, causal
+  predictor that injects the token that historically followed the longest
+  matching suffix of the current context. The ROSA signal is embedded and
+  added to the input representation of each token.
+- **Factorised embedding/head:** 128‑dimensional embedding space, projected
+  to a 512‑dimensional model space, with tied weights.
+- **Context length:** 1024 tokens. No positional embeddings are used, making
+  the context window “free” in terms of parameters.
+- **Tokenizer:** GPT‑2 tokenizer extended with the special tokens `<|user|>`
+  and `<|assistant|>`.
+
+## Uses
+
+### Direct Use
+
+FWKV-ROSA is intended for **research on efficient language models** and for
+**educational demonstrations** of recurrent architectures. You can chat with
+it in a multi‑turn setting using the template above.
+
+### Out‑of‑Scope Use
 
-## Performance
+- This model is **not** suitable for any production or safety‑critical
+  application.
+- It has not been aligned with RLHF or other safety methods and may generate
+  inappropriate or harmful content.
+- The limited size and training data mean it cannot be relied upon for factual
+  knowledge or reasoning.
 
-| Metric | Value |
-|--------|-------|
-| **Validation perplexity** | 67.78 |
-| **Validation loss** | 4.216 |
+### Bias, Risks, and Limitations
 
-*Note:* With only 56M parameters and a fairly simple linear recurrence, the model is intended as a research experiment. It can generate coherent multi‑turn conversations but may occasionally produce repetitive or factual errors.
+- Trained on a relatively small synthetic dataset, the model can produce
+  repetitive or nonsensical output.
+- The ROSA copy mechanism may occasionally copy large chunks of the user’s
+  prompt verbatim.
+- Biases present in the original UltraChat data are likely reflected in the
+  model’s responses.
 
-## How to use
+## How to Get Started
 
-### Load directly with Transformers
+The model relies on a custom architecture. To load it, you must provide the
+`modeling_fwkv.py` file (found in the repository) and trust the remote code:
 
 ```python
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model = AutoModelForCausalLM.from_pretrained("FlameF0X/FWKV-ROSA", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    "FlameF0X/FWKV-ROSA",
+    trust_remote_code=True
+)
 tokenizer = AutoTokenizer.from_pretrained("FlameF0X/FWKV-ROSA")
 ```
 
-**Chat template:**
+Then format your prompts exactly with the chat tokens:
 
+```python
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device).eval()
+
+prompt = "<|user|> What is the capital of France?\n<|assistant|>"
+input_ids = tokenizer.encode(prompt)
+# ROSA IDs must be computed – you can import `rosa` from modeling_fwkv
+from modeling_fwkv import rosa
+rosa_ids = torch.tensor([rosa(input_ids)], device=device)
+out = model(input_ids=torch.tensor([input_ids], device=device),
+            rosa_ids=rosa_ids, use_cache=True)
+# Continue autoregressive sampling…
 ```
-<|user|> Your prompt here
-<|assistant|> The model reply <eos>
-```
 
-Make sure to use the special tokens as shown above. The model was trained to generate after the `<|assistant|>` tag, and the loss is only computed on assistant tokens.
+For a fully working chat demo, see the Gradio app provided in the repository.
+
+## Training Details
+
+### Dataset
 
-## Training details
+- **Name:** [UltraChat 200k](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k)
+- **Splits:** `train_sft` (20,000 examples), `test_sft` (1,000 examples)
+- **Format:** multi‑turn conversations; only assistant tokens contribute to the loss.
+
+### Training Procedure
 
 | Hyperparameter | Value |
 |----------------|-------|
-| Layers | 14 |
-| Model dim | 512 |
-| Embedding dim | 128 |
+| Architecture   | 14 FWKV blocks, d_model=512, d_emb=128 |
 | FFN multiplier | 4 |
-| WKV decay floor | 0.1 |
-| Batch size (per step) | 8 |
+| WKV decay floor| 0.1 |
+| Batch size (per GPU) | 8 |
 | Gradient accumulation | 4 |
-| Effective batch size | 32 |
-| Learning rate | 3e‑4 (cosine schedule) |
-| Weight decay | 0.1 |
+| Effective batch size  | 32 |
+| Learning rate  | 3×10⁻⁴ (cosine schedule) |
+| Weight decay   | 0.1 |
 | Gradient clipping | 1.0 |
-| Epochs | 2 |
-| Optimizer | AdamW (fused) |
-
-Training was done on a single T4 GPU (Kaggle, 2x T4 available but only one used). Gradient checkpointing and bfloat16 autocast were essential to fit the 1024‑token sequences within 15 GB of VRAM. The parallel scan over time is exact (not an approximation) and fully vectorised, yielding a training speed of ~80K tokens/second.
+| Optimizer      | AdamW (fused) |
+| Precision      | bfloat16 mixed |
+| Epochs         | 2 |
+| Hardware       | 1× NVIDIA T4 (15 GB) |
+| Training time  | ~2 hours |
+| Speed          | ~80,000 tokens/second |
+
+Gradient checkpointing was enabled to fit the 1024‑token sequences in memory.
+
+## Evaluation
+
+| Metric                | Value   |
+|-----------------------|---------|
+| Validation loss       | 4.216   |
+| Validation perplexity | 67.78   |
+
+The perplexity is relatively high due to the small model size and limited
+training data. It is comparable to other similarly‑sized recurrent LMs on
+UltraChat.
+
+## Environmental Impact
+
+The training ran for about **2 hours** on a single **NVIDIA T4** GPU
+(maximum power draw ~70 W), resulting in an estimated **0.14 kWh** of
+electricity consumption and approximately **0.06 kg CO₂eq** (assuming a
+grid carbon intensity of 0.4 kg/kWh). This is a negligible footprint.
+
+## Technical Specifications
+
+- **Model type:** Recurrent neural network (linear RNN)
+- **Parameters:** 56.2 million
+  - Backbone (FWKV blocks + embeddings): ~50M
+  - ROSA embedding: ~6.2M
+- **Checkpoint format:** PyTorch `safetensors`
+- **Required files in the repo:**
+  - `config.json`
+  - `model.safetensors` (or `pytorch_model.bin`)
+  - `modeling_fwkv.py`
+  - `tokenizer.json` / `vocab.json` / `merges.txt`
+- **Auto‑mapping:** The `config.json` includes
+  `"auto_map": { "AutoModelForCausalLM": "modeling_fwkv.FWKVLanguageModel" }`,
+  so loading with `trust_remote_code=True` will automatically locate the
+  correct class.
+
+## Citation
+
+If you use FWKV-ROSA in your research, please cite it as:
+
+```bibtex
+@misc{fwkv-rosa,
+  author = {FlameF0X},
+  title = {FWKV-ROSA: A 56M Recurrent Chat LM with RWKV-style Decay and ROSA Copy Signal},
+  year = {2025},
+  howpublished = {\url{https://huggingface.co/FlameF0X/FWKV-ROSA}},
+}
+```
 
-## Limitations & Bias
+## Additional Information
 
-- This is a small model trained on a curated synthetic dialogue dataset. It may not generalise well to factual knowledge or long‑form reasoning.
-- The ROSA copy mechanism helps with repeating patterns but can sometimes cause the model to inappropriately copy parts of the user’s message.
-- No alignment or safety filtering was performed beyond what’s implicit in the dataset.
+This model was built as an experiment to test the combination of a simple
+leaky integrator with the ROSA copy signal on a small, clean conversational
+dataset. It demonstrates that a pure linear RNN can learn to produce
+coherent multi‑turn dialogue without any attention mechanisms. Feedback and
+contributions are welcome!