<h1 align="center">Neurocache</h1>
<h3 align="center">
A library for augmenting language models with external caching mechanisms
</h3>
<a href="https://github.com/alisafaya/neurocache/releases">
<img alt="GitHub release" src="https://img.shields.io/github/release/alisafaya/neurocache.svg">
</a>
## Requirements
* Python 3.6+
* PyTorch 1.13.0+
* Transformers 4.25.0+
## Installation
```bash
pip install neurocache
```
## Getting started
```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from neurocache import (
NeurocacheModelForCausalLM,
OnDeviceCacheConfig,
)
model_name = "facebook/opt-350m"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
cache_layer_idx = model.config.num_hidden_layers - 5
config = OnDeviceCacheConfig(
cache_layers=[cache_layer_idx, cache_layer_idx + 3],
attention_layers=list(range(cache_layer_idx, model.config.num_hidden_layers)),
compression_factor=8,
topk=8,
)
model = NeurocacheModelForCausalLM(model, config)
input_text = ["Hello, my dog is cute", "Hello, my cat is cute"]
tokenized_input = tokenizer(input_text, return_tensors="pt")
tokenized_input["start_of_sequence"] = torch.tensor([0, 1]).bool()
outputs = model(**tokenized_input)
```
## Supported model types
```
from neurocache.utils import NEUROCACHE_SUPPORTED_MODELS
print(NEUROCACHE_SUPPORTED_MODELS)
[
"opt",
"llama",
"mistral",
"gptj",
]
```
Raw data
{
"_id": null,
"home_page": "https://github.com/alisafaya/neurocache",
"name": "neurocache",
"maintainer": "",
"docs_url": null,
"requires_python": ">=3.8.0",
"maintainer_email": "",
"keywords": "deep learning",
"author": "Ali Safaya",
"author_email": "alisafaya@gmail.com",
"download_url": "https://files.pythonhosted.org/packages/75/f6/d47f2a3950bc39359cb943fd6f871c3e47f2d8049b39c62ddca9ad17d9f9/neurocache-0.0.1.tar.gz",
"platform": null,
"description": "\n<h1 align=\"center\">Neurocache</h1>\n<h3 align=\"center\">\n A library for augmenting language models with external caching mechanisms \n</h3>\n\n<a href=\"https://github.com/alisafaya/neurocache/releases\">\n <img alt=\"GitHub release\" src=\"https://img.shields.io/github/release/alisafaya/neurocache.svg\">\n</a>\n\n## Requirements\n\n* Python 3.6+\n* PyTorch 1.13.0+\n* Transformers 4.25.0+\n\n## Installation\n\n```bash\npip install neurocache\n```\n\n## Getting started\n\n```python\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nfrom neurocache import (\n NeurocacheModelForCausalLM,\n OnDeviceCacheConfig,\n)\n\nmodel_name = \"facebook/opt-350m\"\n\nmodel = AutoModelForCausalLM.from_pretrained(model_name)\ntokenizer = AutoTokenizer.from_pretrained(model_name)\n\ncache_layer_idx = model.config.num_hidden_layers - 5\n\nconfig = OnDeviceCacheConfig(\n cache_layers=[cache_layer_idx, cache_layer_idx + 3],\n attention_layers=list(range(cache_layer_idx, model.config.num_hidden_layers)),\n compression_factor=8,\n topk=8,\n)\n\nmodel = NeurocacheModelForCausalLM(model, config)\n\ninput_text = [\"Hello, my dog is cute\", \"Hello, my cat is cute\"]\ntokenized_input = tokenizer(input_text, return_tensors=\"pt\")\ntokenized_input[\"start_of_sequence\"] = torch.tensor([0, 1]).bool()\n\noutputs = model(**tokenized_input)\n```\n\n## Supported model types\n\n```\nfrom neurocache.utils import NEUROCACHE_SUPPORTED_MODELS\nprint(NEUROCACHE_SUPPORTED_MODELS)\n\n[\n \"opt\",\n \"llama\",\n \"mistral\",\n \"gptj\",\n]\n```\n",
"bugtrack_url": null,
"license": "Apache",
"summary": "Neurocache: A library for augmenting language models with external memory.",
"version": "0.0.1",
"project_urls": {
"Homepage": "https://github.com/alisafaya/neurocache"
},
"split_keywords": [
"deep",
"learning"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "af8c685e65a2f7afca86992b738b0065077aa892a616ce71a749d8be8a8cc9c3",
"md5": "0e3f4c837d302cea2cdf22eaa77462e0",
"sha256": "30eaed454d061804ef63f005fe99d0427c407475eac502b03adf8ecc75e94b40"
},
"downloads": -1,
"filename": "neurocache-0.0.1-py3-none-any.whl",
"has_sig": false,
"md5_digest": "0e3f4c837d302cea2cdf22eaa77462e0",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.8.0",
"size": 29485,
"upload_time": "2023-11-06T08:11:57",
"upload_time_iso_8601": "2023-11-06T08:11:57.224994Z",
"url": "https://files.pythonhosted.org/packages/af/8c/685e65a2f7afca86992b738b0065077aa892a616ce71a749d8be8a8cc9c3/neurocache-0.0.1-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "75f6d47f2a3950bc39359cb943fd6f871c3e47f2d8049b39c62ddca9ad17d9f9",
"md5": "b74085eb37b82c523a22f9281201b994",
"sha256": "a7520ad223a1f03a48132adcd189e9a3a8f0ad28363f75395ae52adfb23f55ec"
},
"downloads": -1,
"filename": "neurocache-0.0.1.tar.gz",
"has_sig": false,
"md5_digest": "b74085eb37b82c523a22f9281201b994",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.8.0",
"size": 26164,
"upload_time": "2023-11-06T08:11:59",
"upload_time_iso_8601": "2023-11-06T08:11:59.027342Z",
"url": "https://files.pythonhosted.org/packages/75/f6/d47f2a3950bc39359cb943fd6f871c3e47f2d8049b39c62ddca9ad17d9f9/neurocache-0.0.1.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2023-11-06 08:11:59",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "alisafaya",
"github_project": "neurocache",
"travis_ci": false,
"coveralls": false,
"github_actions": false,
"lcname": "neurocache"
}