This package is to provide a codebase to easily fine-tune some open-source LLMs and reproduce the results.

So far, the codebase supports

• Fine-tuning algorithms offered by the PEFT package
  • LoRA
• Open-source LLM families from Hugging Face
  • Llama2
  • Mistral
  • Falcon
  • Flan-T5
  • Gemma
  • Pythia
• In-context learning

For detailed instruction, please refer to the Notion site: LLM Fine-tuning

---

A Simple Example

Here is a simple example of using the package for fine-tuning or prediction on the SST2 dataset. You can also find this demo code in demo.py.

import fire

from lora import Llama_Lora


def main(
    task: str = "eval",
	base_model: str = "meta-llama/Llama-2-7b-hf",
):
    m = Llama_Lora(
        base_model=base_model,
    )
    if task == "train":
        m.train(
            train_file = "data/sst2/train.json",
            val_file = "data/sst2/val.json",
            output_dir = "./ckp_sst_llama2_lora",
            train_batch_size = 32,
            num_epochs = 1,
        )
    elif task == "eval":
        m.predict(
            input_file = "data/sst2/val.json",
            max_new_tokens = 32,
            verbose = True,
        )
    else:
        raise ValueError(f"Unrecognized task: {task}")


if __name__ == "__main__":
    fire.Fire(main)

For prediction, please use

python demo.py

For fine-tuning, please use

python demo.py --task train

---

Contributors

• Yangfeng Ji
• Wanyu Du
• Aidan San
• Caroline Gihlstorf
• Zhengguang Wang

In addition, this package has been tested by many members from the UVA ILP group via their research projects.

This project was inspired by the Alpaca-LoRA project, from which we learned a lot!