fix md/typos

README.md

@@ -6,12 +6,11 @@ pipeline_tag: text-generation
 tags:
 - text-generation-inference
 ---
-<p align="center" width="100%"><img src="https://raw.githubusercontent.com/CStanKonrad/long_llama/main/assets/longllama.png" alt="LongLLaMA" style="width: 50%;  display: block; margin: auto;"></p>
-
 # LongLLaMA: Focused Transformer Training for Context Scaling
+[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/CStanKonrad/long_llama/blob/main/long_llama_colab.ipynb)
 
 
-[Colab](https://colab.research.google.com/github/CStanKonrad/long_llama/blob/main/long_llama_colab.ipynb) | [TLDR](#TLDR) | [Overview](#Overview) | [Usage](#Usage) | [LongLLaMA performance](#LongLLaMA-performance) | [Authors](#Authors) | [Citation](#Citation) | [License](License) | [Acknowledgments](#Acknowledgments)
+[TLDR](#TLDR) | [Overview](#Overview) | [Usage](#Usage) | [LongLLaMA performance](#LongLLaMA-performance) | [Authors](#Authors) | [Citation](#Citation) | [License](License) | [Acknowledgments](#Acknowledgments)
 
 ## TLDR
 This repository contains the research preview of **LongLLaMA, a large language model capable of handling long contexts of 256k tokens or even more**. 
@@ -24,7 +23,7 @@ LongLLaMA is built upon the foundation of [OpenLLaMA](https://github.com/openlm-
 
 
 **LongLLaMA** is an [OpenLLaMA](https://github.com/openlm-research/open_llama) model finetuned with the FoT method,
-with three layers used for context extension. Crucially, LongLLama is able to extrapolate much beyond the context length seen in training: $8k$. E.g., in the key retrieval task, it can handle inputs of length $256k$.
+with three layers used for context extension. **Crucially, LongLLama is able to extrapolate much beyond the context length seen in training: 8k. E.g., in the key retrieval task, it can handle inputs of length 256k**.
 
 <center>
 
@@ -66,7 +65,7 @@ prompt = "My name is Julien and I like to"
 input_ids = tokenizer(prompt, return_tensors="pt").input_ids
 outputs = model(input_ids=input_ids)
 ```
-During the model call, one can provide the parameter `last_context_length` (default $1024$), which specifies the number of tokens left in the last context window. Tuning this parameter can improve generation as the first layers do not have access to memory. See details in [How LongLLaMA handles long inputs](#How-LongLLaMA-handles-long-inputs).
+During the model call, one can provide the parameter `last_context_length` (default 1024), which specifies the number of tokens left in the last context window. Tuning this parameter can improve generation as the first layers do not have access to memory. See details in [How LongLLaMA handles long inputs](#How-LongLLaMA-handles-long-inputs).
 
 ```python
 generation_output = model.generate(
@@ -85,7 +84,7 @@ LongLLaMA has several other parameters:
 * `mem_layers` specifies layers endowed with memory (should be either an empty list or a list of all memory layers specified in the description of the checkpoint).
 * `mem_dtype` allows changing the type of memory cache
 * `mem_attention_grouping` can trade off speed for reduced memory usage. 
-  When equal to `(4, 2048)`, the memory layers will process at most $4*2048$ queries at once ($4$ heads and $2048$ queries for each head).
+  When equal to `(4, 2048)`, the memory layers will process at most 4*2048 queries at once (4 heads and 2048 queries for each head).
 
 ```python
 import torch
@@ -103,7 +102,7 @@ model = AutoModelForCausalLM.from_pretrained(
 
 
 ### Drop-in use with LLaMA code
- LongLLaMA checkpoints can also be used as a drop-in replacement for LLaMA checkpoints in [Hugging Face implementation of LLaMA](https://huggingface.co/docs/transformers/main/model_doc/llama), but in this case, they will be limited to the original context length of $2048$.
+ LongLLaMA checkpoints can also be used as a drop-in replacement for LLaMA checkpoints in [Hugging Face implementation of LLaMA](https://huggingface.co/docs/transformers/main/model_doc/llama), but in this case, they will be limited to the original context length of 2048.
 
 ```python
 from transformers import LlamaTokenizer, LlamaForCausalLM
@@ -115,11 +114,11 @@ model = LlamaForCausalLM.from_pretrained("syzymon/long_llama_3b", torch_dtype=to
 
 
 ### How LongLLaMA handles long inputs
-Inputs over $2048$ tokens are automatically split into windows $w_1, \ldots, w_m$. The first $m-2$ windows contain $2048$ tokens each, $w_{m-1}$ has no more than $2048$ tokens, and $w_m$ contains the number of tokens specified by `last_context_length`. The model processes the windows one by one extending the memory cache after each. If `use_cache` is `True`, the last window will not be loaded to the memory cache but to the local (generation) cache.
+Inputs over 2048 tokens are automatically split into windows w_1, \ldots, w_m. The first m-2 windows contain 2048 tokens each, w_{m-1} has no more than 2048 tokens, and w_m contains the number of tokens specified by `last_context_length`. The model processes the windows one by one extending the memory cache after each. If `use_cache` is `True`, the last window will not be loaded to the memory cache but to the local (generation) cache.
 
-The memory cache stores $(key, value)$ pairs for each head of the specified memory layers `mem_layers`. In addition to this, it stores attention masks. 
+The memory cache stores (key, value) pairs for each head of the specified memory layers `mem_layers`. In addition to this, it stores attention masks. 
 
-If `use_cache=True` (which is the case in generation), LongLLaMA will use two caches: the memory cache for the specified layers and the local (generation) cache for all layers. When the local cache exceeds $2048$ elements, its content is moved to the memory cache for the memory layers.
+If `use_cache=True` (which is the case in generation), LongLLaMA will use two caches: the memory cache for the specified layers and the local (generation) cache for all layers. When the local cache exceeds 2048 elements, its content is moved to the memory cache for the memory layers.
 
 For simplicity, context extension is realized with a memory cache and full attention in this repo. Replacing this simple mechanism with a KNN search over an external database is possible with systems like [Faiss](https://github.com/facebookresearch/faiss). This potentially would enable further context length scaling. We leave this as a future work.
 
@@ -139,10 +138,10 @@ Our LongLLaMA 3B model also shows improvements when using long context on two do
 
 | Context/Dataset | TREC  | WebQS |
 | --- | --- | --- |
-| $2K$ | 67.0 |  21.2 |
-| $4K$ | 71.6 | 21.4 |
-| $6K$ | 72.9 | 22.2 |
-| $8K$ | **73.3** | **22.4** |
+| 2K | 67.0 |  21.2 |
+| 4K | 71.6 | 21.4 |
+| 6K | 72.9 | 22.2 |
+| 8K | **73.3** | **22.4** |
 
 </center>