🎏 pykoi

Let AI Finetune your Private LLMs

pykoi supercharges your ability to finetune your models

pykoi is an open-source python library to finetune LLMs with RLAIF. Pykoi provides a unified interface including RLHF/RLAIF data and feedback collection, finetuning with reinforcement learning and reward modeling, and LLM comparisons.

Sharable UI for Data & Feedback Collection

Add a chatbot UI with (optional) feedback options to any LLM. The data and interactions will automatically be collected to a database. Currently, two feedback options are supported, vote and rank. Explore how users are using your chatbot, collect feedback, or feed the data directly into RLHF (see below).

1import pykoi as pk
2
3# assume you have some model, endpoint, or api
4# this can be openai, hugggingface, bedrock, claude, etc.
5model = your_model
6
7# give model chat ui with vote feedback
8chatbot = pk.Chatbot(model, feedback="vote")
9# Create sharable link to the application
10app = pk.Application(share=True, username="", password="")
11app.add_component(chatbot)
12app.run()

Easily use RLHF on your own models

Once you've collected enough data, running Reinforcement Learning with Human Feedback (RLHF) is just a few lines of python. Spin up Supervised Finetuning (SFT), Reward Finetuning, and Reinforcement Learning (RL) for an end-to-end RLHF solution. (In the future we will provide our own hosting to make the process even easier!)

1import pykoi as pk
2
3# get data from local database
4qa_database = pk.QuestionAnswerDatabase()
5my_data_pd = qa_database.retrieve_all_question_answers_as_pandas()
6
7# run supervised finetuning
8config = pk.RLHFConfig(base_model_path="meta-llama/Llama-2-7b-hf",
9                          dataset_type="local_db")
10rlhf_step1_sft = pk.SupervisedFinetuning(config)
11rlhf_step1_sft.train_and_save("./models/rlhf_step1_sft")

LLM Comparisons

Easily compare multiple models to each other. Rank the outputs relative to each other, and visualize the results.

1import pykoi as pk
2
3# assume you have some model, endpoint, or api
4model_1 = your_first_model
5model_2 = your_second_model
6model_3 = your_third_model
7model_4 = your_fourth_model
8
9model_array = [
10    model_1,
11    model_2,
12    model_3,
13    model_4
14]
15
16# give interactive chat ui with feedback
17chatbot = pk.Compare(models=model_array)
18
19# Create sharable link to the application
20app = pk.Application(share=False)
21app.add_component(chatbot)
22app.run()