Merge pull request #94 from CambioML/jojo-branch

RAG UI updates
CambioML · Nov 23, 2023 · 61593a5 · 61593a5
2 parents 9b1fcac + 8bad17b
commit 61593a5
Show file tree

Hide file tree

Showing 11 changed files with 229 additions and 99 deletions.
diff --git a/example/retrieval_qa/retrieval_qa_huggingface_demo.py b/example/retrieval_qa/retrieval_qa_huggingface_demo.py
@@ -22,8 +22,7 @@
         # "mistralai/Mistral-7B-v0.1"
         # "databricks/dolly-v2-3b"
 
-RETRIEVAL_MODEL = os.getenv("RETRIEVAL_MODEL")
-
+RETRIEVAL_MODEL = os.getenv("RETRIEVAL_MODEL", default="mistralai/Mistral-7B-v0.1")
 
 def main(**kwargs):
     os.environ["DOC_PATH"] = os.path.join(os.getcwd(), "temp/docs")

diff --git a/pykoi/frontend/dist/assets/index-33aea1b9.css b/pykoi/frontend/dist/assets/index-33aea1b9.css
diff --git a/pykoi/frontend/dist/assets/index-3b11fb1e.js b/pykoi/frontend/dist/assets/index-3b11fb1e.js
diff --git a/pykoi/frontend/dist/assets/index-8dfc5bce.js b/pykoi/frontend/dist/assets/index-8dfc5bce.js
diff --git a/pykoi/frontend/dist/assets/index-fc934c27.css b/pykoi/frontend/dist/assets/index-fc934c27.css
diff --git a/pykoi/frontend/dist/index.html b/pykoi/frontend/dist/index.html
@@ -5,8 +5,8 @@
     <link rel="icon" href="data:image/svg+xml,<svg xmlns=%22http://www.w3.org/2000/svg%22 viewBox=%220 0 100 100%22><text y=%22.9em%22 font-size=%2290%22>🎏</text></svg>">
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
     <title>pykoi</title>
-    <script type="module" crossorigin src="/assets/index-3b11fb1e.js"></script>
-    <link rel="stylesheet" href="/assets/index-fc934c27.css">
+    <script type="module" crossorigin src="/assets/index-8dfc5bce.js"></script>
+    <link rel="stylesheet" href="/assets/index-33aea1b9.css">
   </head>
   <body>
     <div id="app"></div>

diff --git a/pykoi/frontend/src/lib/Chatbots/Components/Dropdown.svelte b/pykoi/frontend/src/lib/Chatbots/Components/Dropdown.svelte
@@ -47,7 +47,6 @@
 
   function handleClickOutside(e) {
     e.preventDefault();
-    console.log('click outside');
     expanded = false;
   }
 
@@ -121,7 +120,7 @@
 
   .dropdown-content {
     position: absolute; /* has to be abs to prevent document overflow */
-    top: 100%;
+    bottom: 100%;
     left: 0;
     width: 100%;
     border: 1px #dadada solid;

diff --git a/pykoi/frontend/src/lib/Chatbots/RAGChatbot.svelte b/pykoi/frontend/src/lib/Chatbots/RAGChatbot.svelte
@@ -33,9 +33,15 @@
     },
   ];
 
+  async function loadRetrievalFilesAndSelect() {
+    await loadRetrievalFiles();
+    $checkedDocs = new Set($uploadedFiles.map((doc) => doc.name));
+    checkedDocs.set(new Set($checkedDocs));
+  }
+
   onMount(() => {
     getDataFromDB();
-    loadRetrievalFiles();
+    loadRetrievalFilesAndSelect();
   });
 
   async function loadRetrievalFiles() {
@@ -237,7 +243,7 @@
     height: 100vh;
     display: grid;
     grid-template-columns: 100%;
-    grid-template-rows: 65% 35%;
+    grid-template-rows: 80% 20%;
   }
 
   .message {

diff --git a/pykoi/frontend/src/lib/RAG/RAG.svelte b/pykoi/frontend/src/lib/RAG/RAG.svelte
@@ -4,13 +4,32 @@
   import { onMount } from "svelte";
   import { uploadedFiles, projections } from "../store.js";
   import CloudArrowUp from "../../assets/CloudArrowUp.svelte";
+  import { slide } from "svelte/transition";
+
+  const UPLOAD_STATES = {
+    WAITING: "waiting",
+    IN_PROGRESS: "in-progress",
+    DONE: "done",
+  };
 
   let selectedFiles = [];
-  let indexed = false;
-  let indexing = false;
+  let uploadState = UPLOAD_STATES.WAITING;
+  let indexState = UPLOAD_STATES.WAITING;
+  let loadState = UPLOAD_STATES.WAITING;
+  let embedState = UPLOAD_STATES.WAITING;
+
+  function resetStates() {
+    uploadState = UPLOAD_STATES.WAITING;
+    indexState = UPLOAD_STATES.WAITING;
+    loadState = UPLOAD_STATES.WAITING;
+    embedState = UPLOAD_STATES.WAITING;
+  }
+
   async function handleFileChange(event) {
     event.preventDefault();
-    let selectedFiles = [];
+    resetStates();
+    uploadState = UPLOAD_STATES.IN_PROGRESS;
+    selectedFiles = [];
     if (event.dataTransfer) {
       if (event.dataTransfer.items) {
         // Use DataTransferItemList interface to access the file(s)
@@ -39,12 +58,18 @@
       method: "POST",
       body: formData,
     });
+    const data = await response.json();
+    console.log("Upload complete! Response:", data);
+    uploadState = UPLOAD_STATES.DONE;
     indexFiles();
     loadServerData();
     getEmbeddings();
   }
 
   async function loadServerData() {
+    if (indexState === UPLOAD_STATES.IN_PROGRESS) {
+      loadState = UPLOAD_STATES.IN_PROGRESS;
+    }
     const response = await fetch("/retrieval/file/get");
     const data = await response.json();
     // Transform the received data
@@ -56,46 +81,57 @@
       };
     });
     $uploadedFiles = [...filesData];
+    if (loadState === UPLOAD_STATES.IN_PROGRESS) {
+      loadState = UPLOAD_STATES.DONE;
+    }
   }
 
   async function indexFiles() {
     console.log("index!");
-    indexing = true;
+    indexState = UPLOAD_STATES.IN_PROGRESS;
     const response = await fetch("/retrieval/vector_db/index", {
       method: "POST",
     });
     const data = await response.json();
-    indexed = true;
-    indexing = false;
+    indexState = UPLOAD_STATES.DONE;
   }
 
   async function getEmbeddings() {
+    embedState = UPLOAD_STATES.IN_PROGRESS;
     console.log("getting embeddings...");
     const response = await fetch("/retrieval/vector_db/get");
     const embeddingData = await response.json();
     console.log("embeddingData", embeddingData);
     $projections = embeddingData;
+    embedState = UPLOAD_STATES.DONE;
   }
 
   function dragOverHandler(event) {
-    console.log("File(s) in drop zone");
-    // Prevent default behavior (Prevent file from being opened)
     event.preventDefault();
   }
 
+  function getSelectedFileNames() {
+    let fileNameStr = "";
+    for (let i = 0; i < selectedFiles.length; i++) {
+      fileNameStr += selectedFiles[i].name + ", ";
+    }
+    return fileNameStr.slice(0, -2);
+  }
+
   onMount(() => {
     loadServerData();
   });
 
-  let dotState = 0;
-
-  // Set an interval to periodically change the number of dots
-  setInterval(() => {
-    dotState = (dotState + 1) % 4;
-  }, 200);
-
-  // Use a reactive statement to create the string with the correct number of dots
-  $: dots = "Indexing" + ".".repeat(dotState);
+  $: {
+    if (
+      uploadState === UPLOAD_STATES.DONE &&
+      indexState === UPLOAD_STATES.DONE &&
+      loadState === UPLOAD_STATES.DONE &&
+      embedState === UPLOAD_STATES.DONE
+    ) {
+      setTimeout(resetStates, 3000);
+    }
+  }
 </script>
 
 <div class="data-grid">
@@ -117,11 +153,19 @@
             Drag and drop files here
           </div>
         </div>
-        {#if indexing && !indexed}
-          <p>{dots}</p>
-        {/if}
-        {#if indexed}
-          <p>Data Successfully indexed!</p>
+        {#if uploadState !== UPLOAD_STATES.WAITING}
+          <div class="processing-container" transition:slide={{ duration: 250 }}>
+            <div>File{selectedFiles.length > 1 ? `s (${selectedFiles.length})` : ""}</div>
+            <div class="processing-files">
+              {getSelectedFileNames()}
+            </div>
+          </div>
+          <div class="upload-status" transition:slide={{ duration: 250 }}>
+            <div class={`loading load-left ${uploadState}`}>Upload</div>
+            <div class={`loading ${indexState}`}>Index</div>
+            <div class={`loading ${loadState}`}>Load</div>
+            <div class={`loading load-right ${embedState}`}>Embed</div>
+          </div>
         {/if}
         <p>Currently pdf, txt, and md are supported.</p>
       </div>
@@ -182,19 +226,93 @@
     flex-direction: column;
     justify-content: center;
     align-items: center;
+    width: 100%;
     height: 100%;
     margin: auto;
     padding: 20px;
     border: 1px solid #333;
     box-sizing: border-box;
+    overflow: hidden;
   }
   .upload-files-container {
     display: flex;
     flex-direction: column;
     gap: 10px;
     width: 90%;
   }
+  .processing-container {
+    color: grey;
+    display: flex;
+    gap: 4px;
+    font-size: small;
+  }
+
+  .processing-files {
+    margin: 0;
+    max-width: 280px;
+    max-height: 2em;
+    overflow-x: auto;
+    overflow-y: hidden;
+    white-space: nowrap;
+    border: 1px solid var(--grey);
+    border-radius: 0.1em;
+    padding: 0.1em 0.5em;
+  }
 
+  .upload-status {
+    display: flex;
+    gap: 3px;
+    justify-content: center;
+    align-items: center;
+    height: 100%;
+    border-radius: 500px;
+    padding: 0 10px;
+  }
+
+  @keyframes color {
+    0% {
+      background-color: var(--yellow);
+    }
+    50% {
+      background-color: var(--lightGrey);
+    }
+    100% {
+      background-color: var(--yellow);
+    }
+  }
+
+  .loading {
+    display: flex;
+    justify-content: center;
+    align-items: center;
+    height: 100%;
+    padding: 0 10px;
+    border-radius: 2px;
+  }
+  .waiting {
+    background-color: var(--lightGrey);
+    color: var(--grey);
+  }
+  .in-progress {
+    background-color: var(--yellow);
+    animation-name: color;
+    animation-duration: 1s;
+    animation-iteration-count: infinite;
+  }
+
+  .done {
+    background-color: var(--green);
+    color: var(--lightGrey);
+  }
+  .load-left {
+    border-top-left-radius: 0.5em;
+    border-bottom-left-radius: 0.5em;
+  }
+
+  .load-right {
+    border-top-right-radius: 0.5em;
+    border-bottom-right-radius: 0.5em;
+  }
   p {
     margin: 0;
   }

diff --git a/pykoi/retrieval/llm/huggingface.py b/pykoi/retrieval/llm/huggingface.py
@@ -7,9 +7,13 @@
 
 from pykoi.retrieval.llm.abs_llm import AbsLlm
 from pykoi.retrieval.vectordb.abs_vectordb import AbsVectorDb
+from dotenv import load_dotenv
 
-MIN_DOCS = 2
+# NOTE: Configure your MIN_DOCS as RAG_NUM_SOURCES in .env file.
+# Load environment variables from .env file
+load_dotenv()
 
+MIN_DOCS = int(os.getenv("RAG_NUM_SOURCES", default=2))
 
 class HuggingFaceModel(AbsLlm):
     """

diff --git a/pykoi/retrieval/llm/openai.py b/pykoi/retrieval/llm/openai.py
@@ -8,10 +8,14 @@
 
 from pykoi.retrieval.llm.abs_llm import AbsLlm
 from pykoi.retrieval.vectordb.abs_vectordb import AbsVectorDb
+from dotenv import load_dotenv
 
+# NOTE: Configure your MIN_DOCS as RAG_NUM_SOURCES in .env file.
+# Load environment variables from .env file
+load_dotenv()
 
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-MIN_DOCS = 2
+MIN_DOCS = int(os.getenv("RAG_NUM_SOURCES", default=2))
 
 
 class OpenAIModel(AbsLlm):
@@ -26,7 +30,7 @@ def __init__(self, vector_db: AbsVectorDb):
         try:
             self._llm = OpenAI(
                 model_name="gpt-4",
-                temperature=0, 
+                temperature=0,
                 max_tokens=500)
 
             self._vector_db = vector_db.vector_db