add readme

.ipynb_checkpoints/README-checkpoint.md

@@ -1,3 +1,38 @@
 ---
 license: mit
 ---
+
+# [Gated State Space](https://arxiv.org/abs/2206.13947)
+
+This repo contains pretrain model for the gated state space paper. The model has been trained on [C4 dataset](https://huggingface.co/datasets/c4). I have used [Lucidrains' implementation](https://github.com/lucidrains/gated-state-spaces-pytorch) ([commit](https://github.com/lucidrains/gated-state-spaces-pytorch/tree/32cd036e775112cc469e94fa1165fe111393708b)) for the model. I think the main benefit of this model is the ability to scale beyond the training context length. As authors noted in the paper, they trained the model on 4k sequence length but it generalized beyond that length. I have written a blog post on how I started the training [here](https://naxalpha.substack.com/p/devlog-experiment-a2a468-gated-state).
+
+[Here](https://wandb.ai/naxalpha/gated-state-space/reports/Gated-State-Space-Training-v1--VmlldzozMTYzMzY3?accessToken=zy10rrpofi9k7l52aqwiej8bk0ub302rdswfkxmf8y94dt2j6z4kxbca6ar3sc52) are the training logs (report) on W&B. Since the training is still going on, this repo currently contains the checkpoint @ ~200k step. The report contains both the loss and the output results. While the generated text somewhat makes sense, given the loss is still ~4.2, it is not useful out-of-the-box for most tasks. So need to fine-tune it before using it for anything else.
+
+## How to use this.
+
+Since it is not based on [transformers](https://github.com/huggingface/transformers/) library, it is a bit tricky to use the model out of the box. Here are the general steps:
+
+1. `pip install gated-state-spaces-pytorch`
+2. Download the model weights from [here](https://huggingface.co/naxalpha/gated-state-space/raw/main/model.pt).
+3. Download the config from [here](https://huggingface.co/naxalpha/gated-state-space/raw/main/config.json).
+4. Following code to patch the original model:
+```python
+    model = AutoregressiveWrapper(
+        GatedStateSpacesLM(
+            **config
+        ),
+    )
+    model.net.to_logits = nn.Sequential(
+        nn.LayerNorm(f_emb),
+        model.net.to_logits,
+    )
+```
+5. Load the state dict: `model.load_state_dict(torch.load('model.pt'))`
+6. If you want to fine-tune the model, you can freeze the embeddings:
+```python
+    model.net.token_emb.weight.requires_grad_(False)
+    model.net.token_emb.weight.copy_(emb)
+
+    model.net.to_logits[1].weight.requires_grad_(False)
+    model.net.to_logits[1].weight.copy_(emb)
+```

README.md

@@ -1,3 +1,38 @@
 ---
 license: mit
 ---
+
+# [Gated State Space](https://arxiv.org/abs/2206.13947)
+
+This repo contains pretrain model for the gated state space paper. The model has been trained on [C4 dataset](https://huggingface.co/datasets/c4). I have used [Lucidrains' implementation](https://github.com/lucidrains/gated-state-spaces-pytorch) ([commit](https://github.com/lucidrains/gated-state-spaces-pytorch/tree/32cd036e775112cc469e94fa1165fe111393708b)) for the model. I think the main benefit of this model is the ability to scale beyond the training context length. As authors noted in the paper, they trained the model on 4k sequence length but it generalized beyond that length. I have written a blog post on how I started the training [here](https://naxalpha.substack.com/p/devlog-experiment-a2a468-gated-state).
+
+[Here](https://wandb.ai/naxalpha/gated-state-space/reports/Gated-State-Space-Training-v1--VmlldzozMTYzMzY3?accessToken=zy10rrpofi9k7l52aqwiej8bk0ub302rdswfkxmf8y94dt2j6z4kxbca6ar3sc52) are the training logs (report) on W&B. Since the training is still going on, this repo currently contains the checkpoint @ ~200k step. The report contains both the loss and the output results. While the generated text somewhat makes sense, given the loss is still ~4.2, it is not useful out-of-the-box for most tasks. So need to fine-tune it before using it for anything else.
+
+## How to use this.
+
+Since it is not based on [transformers](https://github.com/huggingface/transformers/) library, it is a bit tricky to use the model out of the box. Here are the general steps:
+
+1. `pip install gated-state-spaces-pytorch`
+2. Download the model weights from [here](https://huggingface.co/naxalpha/gated-state-space/raw/main/model.pt).
+3. Download the config from [here](https://huggingface.co/naxalpha/gated-state-space/raw/main/config.json).
+4. Following code to patch the original model:
+```python
+    model = AutoregressiveWrapper(
+        GatedStateSpacesLM(
+            **config
+        ),
+    )
+    model.net.to_logits = nn.Sequential(
+        nn.LayerNorm(f_emb),
+        model.net.to_logits,
+    )
+```
+5. Load the state dict: `model.load_state_dict(torch.load('model.pt'))`
+6. If you want to fine-tune the model, you can freeze the embeddings:
+```python
+    model.net.token_emb.weight.requires_grad_(False)
+    model.net.token_emb.weight.copy_(emb)
+
+    model.net.to_logits[1].weight.requires_grad_(False)
+    model.net.to_logits[1].weight.copy_(emb)
+```