Merge branch 'main' into flex

.gitignore

@@ -176,6 +176,8 @@ rlhf_tests
 test_data
 example/flex/docs
 example/flex/flex.ipynb
+example/flex/flex_sample/
+example/flex/vectordb/
 
 # mac
 .DS_Store
@@ -189,3 +191,6 @@ models/
 
 # retrieval
 temp/
+
+# poetry
+poetry.lock

README.md

@@ -1,75 +1,97 @@
 
-# 🎏 pykoi: Active learning in one unified interface
+# 🎏 pykoi: RLHF/RLAIF in one unified interface
+<p align="center">
+  <a href="/LICENSE"><img alt="License Apache-2.0" src="https://img.shields.io/github/license/cambioml/pykoi?style=flat-square"></a>
+  <a href="https://pypi.org/project/pykoi"><img src="https://img.shields.io/pypi/v/pykoi.svg" alt="pypi_status" /></a>
+  <a href="https://github.com/cambioml/pykoi/graphs/commit-activity"><img alt="Commit activity" src="https://img.shields.io/github/commit-activity/m/cambioml/pykoi?style=flat-square"/></a>
+  <a href="https://join.slack.com/t/cambiomlworkspace/shared_invite/zt-1zes33rmt-20Rag043uvExUaUdvt5_xQ"><img src="https://badgen.net/badge/Join/Community/cyan?icon=slack" alt="Slack" /></a>
+</p>
+
+[pykoi](https://www.cambioml.com/pykoi/) is an open-source python library for improving LLMs with RLHF. We provide a unified interface including RLHF/RLAIF data and feedback collection, finetuning with reinforcement learning and reward modeling, and LLM comparisons. 
+
+## Features
+
+`pykoi` let you easily get real-time user feedback and continuously improve your models. Here are some common applications:
+
+### Sharable UI
+
+Do you want to store your chat history with LLMs from OpenAI, Amazon Bedrock(:woman_technologist: building now :man_technologist:), or Huggingface? With just three lines of code, pykoi lets you to store them locally, ensuring 100% privacy. This includes launching a chatbot UI, automatically saving your chat history in your compute instance (cpu or gpu), and visualizing it on a dashboard. Explore the demos below:
+
+- If you're using a CPU instance, check out [demo_launch_app_cpu.ipynb](https://nbviewer.org/github/CambioML/pykoi/blob/main/example/chatbot/demo_launch_app_cpu_openai.ipynb)
+- If you're using a GPU instance, check out [demo_launch_app_gpu.ipynb](https://nbviewer.org/github/CambioML/pykoi/blob/main/example/chatbot/demo_launch_app_gpu_huggingface.ipynb)
+- Alternatively, read our [blog](https://www.cambioml.com/docs/data_collection_feedback.html) for more information!
+
+![Watch the video](example/image/pykoi_demo_rlaif_data_collection.gif)
+
+
+### Model comparison
+
+Comparing models is a difficult task. `pykoi` makes it easy by allowing one to directly compare the performance of multiple models to each other, with just a few lines of code. If you have multiple language models that you’d like to compare to each other on a set of prompts or via an interactive session, you can use `pk.Compare`. Check out any of the demo below: 
+
+- If you're using a CPU instance, check out [demo_launch_app_cpu.ipynb](https://nbviewer.org/github/CambioML/pykoi/blob/main/example/chatbot/demo_model_comparator_openai.ipynb)
+- If you're using a GPU instance, check out [demo_launch_app_gpu.ipynb](https://nbviewer.org/github/CambioML/pykoi/blob/main/example/chatbot/demo_model_comparator_gpu_huggingface.ipynb)
+- Alternatively, read our [blog](https://www.cambioml.com/docs/model_comparison.html) for more information!
+
+ ![Watch the video](example/image/pykoi_demo_model_comparison.gif)
+
+### RLHF
+
+Reinforcement Learning with Human Feedback (RLHF) is a unique training paradigm that blends reinforcement learning with human-in-the-loop training. The central idea is to use human evaluative feedback to refine a model's decision-making ability and guide the learning process towards desired outcomes. Researchers from [Deepmind](https://www.deepmind.com/blog/learning-through-human-feedback), [OpenAI](https://openai.com/research/learning-from-human-preferences) and [Meta Llama2](https://arxiv.org/pdf/2307.09288.pdf) have all demonstrated that RLHF is a game changer for large language models (LLMs) training.
+
+`pykoi` allows you to easily fine-tune your model on the datasets you've collected via your `pykoi` chat or rank databases. Check our [blog](https://www.cambioml.com/docs/rlhf.html) for detailed instructions on how to use it.
 
-[pykoi](https://www.cambioml.com/pykoi/) is an open-source python library for LLMs providing a unified interface for data & feedback collection, RLHF, and model comparisons.
 
 ## Installation
-To get started with `pykoi`, you can choose to one of following compute options: CPU (e.g. your laptop) or GPU (e.g. EC2).
+To get started with pykoi, you can choose from any of the installation options. The choice should be based on the features you need (e.g., RAG, RLHF or all) and the compute resources you have, such as a CPU (e.g., your laptop) or GPU (e.g., AWS EC2 or SageMaker).
 
-### Option 1: CPU (e.g. your laptop)
-Installation on a CPU is simple if you have conda. If not, install [conda](https://docs.conda.io/projects/conda/en/latest/user-guide/install/index.html) for your operating system.
+### Option 1: RAG (CPU)
+This option allows you to run RAG on a CPU using either the OpenAI API or the Anthropic Claude2 API. Installation of RAG (CPU) is simple if you have conda. If not, install [conda](https://docs.conda.io/projects/conda/en/latest/user-guide/install/index.html) for your operating system.
 
 First, create a conda environment on your terminal using:
 ```
 conda create -n pykoi python=3.10 -y
-conda activate pykoi
+conda activate pykoi  # some OS requires `source activate pykoi`
 ```
 
-Then install `pykoi` and the compatible torch version.
+Then install `pykoi` and the compatible [pytorch based on your os](https://pytorch.org/get-started)
 ```
-pip3 install pykoi
+pip3 install pykoi[rag]
+pip3 install torch 
 ```
 
-### Option 2:  GPU (e.g. EC2)
+### Option 2: RAG (GPU) 
+This option allows you to run RAG on a GPU using an open-source LLM from HuggingFace. Here's a quick [tutorial](#ec2-dev-setup) on setting up an EC2 GPU instance for the installation below.
 
-If you are on EC2, you can launch a GPU instance with the following config:
-- EC2 `g4dn.xlarge` (if you want to run a pretrained model with 7B parameters)
-- Deep Learning AMI PyTorch GPU 2.0.1 (Ubuntu 20.04)
-    <img src="example/image/readme_ec2_ami.jpg" alt="Alt text" width="50%" height="50%"/>
-- EBS: at least 100G
-    <img src="example/image/readme_ec2_storage.png" alt="Alt text" width="50%" height="50%"/>
-
-Next, on your GPU instance terminal, create a conda environment using:
+On your GPU instance terminal, create a conda environment using:
 ```
 conda create -n pykoi python=3.10 -y && source activate pykoi
 ```
 
-Then install `pykoi` and the correlated torch version.
+Then install `pykoi` and [pytorch based on your cuda version](https://pytorch.org/get-started). You can find your CUDA version via `nvcc -V`.
 ```
-pip3 install pykoi && pip3 install torch --index-url https://download.pytorch.org/whl/cu118
-```
-
-
-Congrats you have finished installation, 
+pip3 install pykoi[huggingface]
 
-## How do I use `pykoi`?
-
-`pykoi` is a python interface to unify your ML model development and production. You can easily get real-time user feedback and continuously improve your model.
-
-Here are some common applications:
-
-### Sharable Chatbot UI with Database
-
-Are you concerned about storing your chat history with chatbots from OpenAI, Amazon Bedrock(:woman_technologist: building now :man_technologist:), or Huggingface? With just three lines of code, we enable you to store it locally, ensuring 100% privacy. This includes launching a chatbot UI, automatically saving your chat history, and visualizing it on a dashboard. Explore the demos below:
-
-- If you're using a CPU instance, check out [demo_launch_app_api.ipynb](https://nbviewer.org/github/CambioML/pykoi/blob/main/example/chatbot/demo_launch_app_api.ipynb)
-- If you're using a GPU instance, check out [demo_launch_app_gpu.ipynb](https://nbviewer.org/github/CambioML/pykoi/blob/main/example/chatbot/demo_launch_app_gpu.ipynb)
-- Alternatively, read our [blog](https://www.cambioml.com/docs/data_collection_feedback.html) for more information!
-
-### Model comparison
+# install torch based on cuda (e.g. cu118 means cuda 11.8)
+pip3 install torch --index-url https://download.pytorch.org/whl/cu118
+```
 
-Comparing models is a difficult task. `pykoi` makes it easy by allowing one to directly compare the performance of multiple models to each other, with just 1 line of code. If you have multiple language models that you’d like to compare to each other on a set of prompts or via an interactive session, you can use `pk.Compare`. Check out any of the demo below: 
+### Option 3: RLHF (GPU) 
+This option allows you to train LLM via RLHF on a GPU. Here's a quick [tutorial](#ec2-dev-setup) on setting up an EC2 GPU instance for the installation below.
 
-- If you're using a CPU instance, check out [demo_launch_app_api.ipynb](https://nbviewer.org/github/CambioML/pykoi/blob/main/example/chatbot/demo_model_comparator_openai.ipynb)
-- If you're using a GPU instance, check out [demo_launch_app_gpu.ipynb](https://nbviewer.org/github/CambioML/pykoi/blob/main/example/chatbot/demo_model_comparator_hf.ipynb)
-- Alternatively, read our [blog](https://www.cambioml.com/docs/model_comparison.html) for more information!
-
-### RLHF
+On your GPU instance terminal, create a conda environment using:
+```
+conda create -n pykoi python=3.10 -y && source activate pykoi
+```
 
-Reinforcement Learning with Human Feedback (RLHF) is a unique training paradigm that blends reinforcement learning with human-in-the-loop training. The central idea is to use human evaluative feedback to refine a model's decision-making ability and guide the learning process towards desired outcomes. Researchers from [Deepmind](https://www.deepmind.com/blog/learning-through-human-feedback), [OpenAI](https://openai.com/research/learning-from-human-preferences) and [Meta Llama2](https://arxiv.org/pdf/2307.09288.pdf) have all demonstrated that RLHF is a game changer for large language models (LLMs) training.
+Then install `pykoi` and [pytorch based on your cuda version](https://pytorch.org/get-started). You can find your CUDA version via `nvcc -V`.
+```
+pip3 install pykoi[rlhf]
 
-`pykoi` allows you to easily fine-tune your model on the datasets you've collected via your `pykoi` chat or rank databases. Check our [blog](https://www.cambioml.com/docs/rlhf.html) for detailed instructions on how to use it.
+# install torch based on cuda (e.g. cu118 means cuda 11.8)
+pip3 install torch --index-url https://download.pytorch.org/whl/cu118
+```
 
+Congrats you have finished the installation! 
 
 ## Dev Setup
 If you are interested to contribute to us, here are the preliminary development setup.
@@ -80,13 +102,44 @@ conda create -n pykoi python=3.10
 conda activate pykoi
 cd pykoi
 pip3 install poetry
-poetry install --no-root
 ```
 
+Then, based the feature you need to develop, run one or more installation options below. We recommend install all the options below although it may take ~3 minutes longer.
+
+- Option 1: RAG (CPU) 
+  ```
+  poetry install --no-root --extras rag
+  ```
+- Option 2: RAG (GPU) 
+  ```
+  poetry install --no-root --extras huggingface
+  ```
+- Option 3: RLHF (GPU) 
+  ```
+  poetry install --no-root --extras rlhf
+  ```
+
+Finally, if you are on a GPU, install `pykoi` and [pytorch based on your cuda version](https://pytorch.org/get-started). You can find your CUDA version via `nvcc -V`.
+```
+pip3 install pykoi[huggingface]
+
+# install torch based on cuda (e.g. cu118 means cuda 11.8)
+pip3 install torch --index-url https://download.pytorch.org/whl/cu118
+```
+
+
 ### Frontend Dev Setup
-Frontend:
 ```
 cd pykoi/pykoi/frontend
 npm install
 npm run build
 ```
+
+### EC2 Dev Setup
+If you are on EC2, you can launch a GPU instance with the following config:
+- EC2 `g4dn.xlarge` (if you want to run a pretrained LLM with 7B parameters)
+- Deep Learning AMI PyTorch GPU 2.0.1 (Ubuntu 20.04)
+    <img src="example/image/readme_ec2_ami.jpg" alt="Alt text" width="50%" height="50%"/>
+- EBS: at least 100G
+    <img src="example/image/readme_ec2_storage.png" alt="Alt text" width="50%" height="50%"/>
+

docker/README.md

@@ -0,0 +1,14 @@
+# Overview
+In this folder, we create the different dockerfiles for using pykoi.
+
+1. `pykoi-cpu`: The base image for the cpu-based usage.
+2. `pykoi-cpu-custom`: When you run this docker image, try to modify the `app.py` and mount it when running the docker container.
+
+To run a docker container, we can use the following command:
+```bash
+docker run -dp 5000:5000 -v $(pwd)/app.py:/app/app.py \
+        --name alex_test \
+        pykoi/pykoi-cpu:app
+```
+
+Note that we need to keep the exposed port as 5000 (default value) to make the server work. 

docker/pykoi-cpu-custom/Dockerfile

@@ -0,0 +1,13 @@
+FROM pykoi/pykoi-cpu:base
+
+# set the working directory in the container
+WORKDIR /app
+
+# copy the content
+COPY . /app
+
+# Expose the port on which the app will be running 
+EXPOSE 5000
+
+# For dev
+CMD ["python", "app.py"]

example/chatbot/openai_model_demo.py → docker/pykoi-cpu-custom/app.py

@@ -5,7 +5,7 @@
 # Creating an OpenAI model (requires an OpenAI API key) #
 ##########################################################
 # enter openai api key here
-api_key = ""
+api_key = "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"
 
 # Creating an OpenAI model
 model = pykoi.ModelFactory.create_model(
@@ -25,7 +25,9 @@
 ###########################################################
 # Create the application
 # app = pykoi.Application(debug=False, share=True)
-app = pykoi.Application(debug=False, share=False)
+app = pykoi.Application(
+    debug=False, 
+    share=True)
 app.add_component(chatbot)
 app.add_component(dashboard)
 app.run()

docker/pykoi-cpu/Dockerfile

@@ -0,0 +1,27 @@
+FROM python:3.10
+LABEL maintainer="Pykio Team"
+LABEL repository="Pykoi"
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt update && \
+    apt install -y bash \
+                   build-essential \
+                   git \
+                   git-lfs \
+                   curl \
+                   ca-certificates \
+                   libsndfile1-dev
+
+# make sure to use venv
+RUN python -m venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+
+# pre-install the heavy dependencies (these can later be overridden by the deps from setup.py)
+# follow the instructions here: https://cloud.google.com/tpu/docs/run-in-container#train_a_jax_model_in_a_docker_container
+RUN python -m pip install --no-cache-dir --upgrade pip && \
+    python -m pip install --upgrade --no-cache-dir \
+    pykoi
+
+CMD ["/bin/bash"]
+

docker/pykoi-gpu/Dockerfile

@@ -0,0 +1,31 @@
+FROM python:3.10
+LABEL maintainer="Pykio Team"
+LABEL repository="Pykoi"
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt update && \
+    apt install -y bash \
+                   build-essential \
+                   git \
+                   git-lfs \
+                   curl \
+                   ca-certificates \
+                   libsndfile1-dev \
+                   libgl1
+
+# make sure to use venv
+RUN python3 -m venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+
+# pre-install the heavy dependencies (these can later be overridden by the deps from setup.py)
+RUN python3 -m pip install --no-cache-dir --upgrade pip && \
+    python3 -m pip install --no-cache-dir \
+        torch \
+        torchvision \
+        torchaudio && \
+    python3 -m pip install --no-cache-dir \
+        pykoi
+
+CMD ["/bin/bash"]
+

example/chatbot/chatbot_in_jupyter.ipynb

@@ -16,9 +16,20 @@
     "current_directory = os.getcwd()\n",
     "sys.path.append(\".\")\n",
     "sys.path.append(\"..\")\n",
-    "sys.path.append(\"../..\")\n",
-    "\n",
-    "import pykoi as pk"
+    "sys.path.append(\"../..\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6a907bb3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pykoi import Application\n",
+    "from pykoi.chat import ModelFactory\n",
+    "from pykoi.chat import QuestionAnswerDatabase\n",
+    "from pykoi.component import Chatbot"
    ]
   },
   {
@@ -31,7 +42,7 @@
     "api_key = \"\"\n",
     "\n",
     "# Creating an OpenAI model\n",
-    "model = pk.ModelFactory.create_model(model_source=\"openai\", api_key=api_key)"
+    "model = ModelFactory.create_model(model_source=\"openai\", api_key=api_key)"
    ]
   },
   {
@@ -41,10 +52,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "database = pk.QuestionAnswerDatabase(debug=True)\n",
-    "chatbot = pk.Chatbot(model=model, feedback=\"vote\")\n",
+    "database = QuestionAnswerDatabase(debug=True)\n",
+    "chatbot = Chatbot(model=model, feedback=\"vote\")\n",
     "\n",
-    "app = pk.Application(debug=False, share=False)\n",
+    "app = Application(debug=False, share=False)\n",
     "app.add_component(chatbot)"
    ]
   },
@@ -76,7 +87,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.10.12"
   }
  },
  "nbformat": 4,