Update retrieval demo

README.md

@@ -7,7 +7,7 @@
   <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. 
+[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
 
@@ -26,12 +26,12 @@ Do you want to store your chat history with LLMs from OpenAI, Amazon Bedrock(:wo
 
 ### 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: 
+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
@@ -40,6 +40,19 @@ Reinforcement Learning with Human Feedback (RLHF) is a unique training paradigm
 
 `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.
 
+### RAG
+Implement a Retrieval-Augmented Generation (RAG) Chatbot quickly with `pykoi`. Upload your own documents to create context-aware responses on top of a pretrained LLM, such as from OpenAI or Huggingface.
+
+![Watch the video](example/image/pykoi_demo_rag_chatbot_upload.gif)
+
+Then chat with the RAG Chatbot, where you can select specific sources for the text generation, as well as see which sources the model used to generate the response.
+
+![Watch the video](example/image/pykoi_demo_rag_chatbot_chat.gif)
+
+Finally, you can save a modified response to help improve the model. This is a great way to collect data for RLHF.
+
+![Watch the video](example/image/pykoi_demo_rag_chatbot_modify.gif)
+
 
 ## Installation
 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).
@@ -56,10 +69,10 @@ conda activate pykoi  # some OS requires `source activate pykoi`
 Then install `pykoi` and the compatible [pytorch based on your os](https://pytorch.org/get-started/locally/)
 ```
 pip3 install "pykoi[rag]"
-pip3 install torch 
+pip3 install torch
 ```
 
-### Option 2: RAG (GPU) 
+### 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.
 
 On your GPU instance terminal, create a conda environment using:
@@ -74,7 +87,7 @@ pip3 install "pykoi[rag, huggingface]"
 pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121  # cu121 means cuda 12.1
 ```
 
-### Option 3: RLHF (GPU) 
+### 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.
 
 On your GPU instance terminal, create a conda environment using:
@@ -89,7 +102,7 @@ pip3 install "pykoi[rlhf]"
 pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121  # cu121 means cuda 12.1
 ```
 
-Congrats you have finished the installation! 
+Congrats you have finished the installation!
 
 ## Dev Setup
 If you are interested to contribute to us, here are the preliminary development setup.
@@ -104,15 +117,15 @@ pip3 install poetry
 
 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) 
+- Option 1: RAG (CPU)
   ```
   poetry install --no-root --extras rag
   ```
-- Option 2: RAG (GPU) 
+- Option 2: RAG (GPU)
   ```
   poetry install --no-root --extras "rag huggingface"
   ```
-- Option 3: RLHF (GPU) 
+- Option 3: RLHF (GPU)
   ```
   poetry install --no-root --extras rlhf
   ```

example/retrieval_qa/retrieval_qa_huggingface_demo.py

@@ -3,7 +3,7 @@
 import os
 import argparse
 from pykoi import Application
-from pykoi.chat import QuestionAnswerDatabase
+from pykoi.chat import RAGDatabase
 from pykoi.retrieval import RetrievalFactory
 from pykoi.retrieval import VectorDbFactory
 from pykoi.component import Chatbot, Dashboard, RetrievalQA
@@ -36,9 +36,9 @@ def main(**kwargs):
     )
 
     # retrieval, chatbot, and dashboard pykoi components
-    retriever = RetrievalQA(retrieval_model=retrieval_model, vector_db=vector_db)
-    chatbot = Chatbot(None, feedback="vote", is_retrieval=True)
-    dashboard = Dashboard(QuestionAnswerDatabase())
+    retriever = RetrievalQA(retrieval_model=retrieval_model, vector_db=vector_db, feedback="rag")
+    chatbot = Chatbot(None, feedback="rag", is_retrieval=True)
+    dashboard = Dashboard(RAGDatabase(), feedback="rag")
 
     ############################################################
     # Starting the application and retrieval qa as a component #

example/flex/flex_demo.py → example/retrieval_qa/retrieval_qa_openai_demo.py

@@ -4,7 +4,6 @@
 import argparse
 from dotenv import load_dotenv
 from pykoi import Application
-from pykoi.chat import QuestionAnswerDatabase
 from pykoi.retrieval import RetrievalFactory
 from pykoi.retrieval import VectorDbFactory
 from pykoi.component import Chatbot, Dashboard, RetrievalQA

pykoi/application.py

@@ -644,13 +644,13 @@ async def inference(
             try:
                 print("[/retrieval]: model inference.....", request_body.prompt)
                 component["component"].retrieval_model.re_init(request_body.file_names)
-                output = component[
-                    "component"
-                ].retrieval_model.run_with_return_source_documents(
-                    {"query": request_body.prompt}
-                )
-                print("output", output, output["result"])
-                if output["source_documents"] == []:
+                output = component["component"].retrieval_model.run_with_return_source_documents({"query": request_body.prompt})
+                print('output', output, output["result"])
+                if "source_documents" not in output:
+                    print('no source documents', output)
+                    source = ["N/A"]
+                    source_content = ["N/A"]
+                elif output["source_documents"] == []:
                     source = ["N/A"]
                     source_content = ["N/A"]
                 else: