nano-askllm 0.1.0

Unofficial implementation of the paper 'How to Train Data-Efficient LLMs' arXiv:2402.09668

nano-askllm

Unofficial implementation of the paper 'How to Train Data-Efficient LLMs' arXiv:2402.09668

Installation

pip install nano-askllm

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
from nano_askllm import AskLLM

datapoints = [
  "first datapoint",  # See tests/test_askllm.py for examples
  "second datapoint",
]

model_id = "google/gemma-2b-it"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

prompt_template = "..."  # See tests/test_askllm.py for an example
yes_tokens = ["yes", "Yes", "YES", " yes", " Yes", " YES"]

llm = AskLLM(tokenizer, model, prompt_template=prompt_template, yes_tokens=yes_tokens)
prompts = llm.get_prompts(datapoints)
inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device)
results = llm.ask(inputs)

print(results)

Development

poetry install
poetry run pytest -s     # run pytest once
poetry run -- ptw -- -s  # watch for changes and run pytest

Citation

@misc{sachdeva2024train,
      title={How to Train Data-Efficient LLMs},
      author={Noveen Sachdeva and Benjamin Coleman and Wang-Cheng Kang and Jianmo Ni and Lichan Hong and Ed H. Chi and James Caverlee and Julian McAuley and Derek Zhiyuan Cheng},
      year={2024},
      eprint={2402.09668},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

License

MIT License. See LICENSE for details.