Update README.md

README.md

@@ -1,10 +1,91 @@
+Here’s a proper, detailed README for your model repository. It covers the architecture, training, usage, and performance.
+
+```markdown
 ---
 language: en
 license: apache-2.0
 tags:
   - fwkv
   - rosa
+  - rwkv
+  - linear-transformer
+  - recurrent-neural-network
+  - chat
+  - ultrachat
+  - pytorch
+  - custom-architecture
 ---
-# FWKV-56M + ROSA (chat)
-Factorised tied weights, vectorized decayed-accumulator recurrence, RWKV-8 ROSA copy signal, chat-tuned on HuggingFaceH4/ultrachat_200k.
-**Best val perplexity:** 67.78
+
+# 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.
+
+## Model description
+
+- **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.
+
+## Performance
+
+| Metric | Value |
+|--------|-------|
+| **Validation perplexity** | 67.78 |
+| **Validation loss** | 4.216 |
+
+*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.
+
+## How to use
+
+### Load directly with Transformers
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# The architecture is custom – you need the model definition from the repository.
+# The simplest way is to use the same class definition as in the training script,
+# or rely on the provided `modeling_fwkv.py` (if you upload it) and `trust_remote_code=True`.
+
+model = AutoModelForCausalLM.from_pretrained("FlameF0X/FWKV-ROSA", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("FlameF0X/FWKV-ROSA")
+```
+
+**Chat template:**
+
+```
+<|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.
+
+## Training details
+
+| Hyperparameter | Value |
+|----------------|-------|
+| Layers | 14 |
+| Model dim | 512 |
+| Embedding dim | 128 |
+| FFN multiplier | 4 |
+| WKV decay floor | 0.1 |
+| Batch size (per step) | 8 |
+| Gradient accumulation | 4 |
+| Effective batch size | 32 |
+| Learning rate | 3e‑4 (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.
+
+## Limitations & Bias
+
+- 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.