add sink cache options

README.md

@@ -17,12 +17,17 @@ This implementation should match the `SinkCache` class present in `transformers<
 
 
 ## Model compatibility
+- Decoder-only models
 
 
 ## Additional Arguments
+- `window_length` (`int`, defaults to `256`): The length of the context window.
+- `num_sink_tokens` (`int`, defaults to `4`): The number of sink tokens. See the original paper for more information.
 
 
 ## Output Type changes
+- When `return_dict_in_generate=True`, `output.past_key_values` will be a `SinkCache` instance. `SinkCache` is defined
+in `generate.py`, in this repository.
 
 
 ## Example usage

custom_generate/generate.py

@@ -193,7 +193,7 @@ class SinkCache(Cache):
         return self.key_cache[layer_idx], self.value_cache[layer_idx]
 
 
-def generate(model, **kwargs):
-    past_key_values = SinkCache(window_length=256, num_sink_tokens=4)
+def generate(model, window_length=256, num_sink_tokens=4, **kwargs):
+    past_key_values = SinkCache(window_length=window_length, num_sink_tokens=num_sink_tokens)
     generation_outputs = model.generate(**kwargs, past_key_values=past_key_values, use_cache=True)
     return generation_outputs

custom_generate/requirements.txt

@@ -0,0 +1 @@
+transformers>=4.53.0  # 4.52 results in an infinite loop