Host data on cloudflare

Signed-off-by: Jay Wang <[email protected]>
poloclub · Feb 9, 2024 · 6760029 · 6760029
1 parent e1de323
commit 6760029
Show file tree

Hide file tree

Showing 5 changed files with 259 additions and 19 deletions.
diff --git a/examples/rag-playground/src/components/app/app.css b/examples/rag-playground/src/components/app/app.css
@@ -6,6 +6,8 @@
   flex-direction: column;
   justify-content: flex-start;
   align-items: center;
+
+  --background-color: hsl(206, 75%, 17%);
 }
 
 * {
@@ -39,7 +41,7 @@
   box-sizing: border-box;
   overflow-x: hidden;
   position: relative;
-  background-color: hsl(206, 75%, 17%);
+  background-color: var(--background-color);
 }
 
 .text-left {
@@ -156,3 +158,85 @@
   margin-top: 80px;
   color: #e0e0e0;
 }
+
+button {
+  all: unset;
+}
+
+.app-tabs {
+  display: flex;
+  flex-direction: row;
+  padding: 10px 0 0 0;
+
+  .splitter {
+    border-right: 1px solid var(--gray-400);
+    margin: 0 15px;
+    width: 0px;
+  }
+
+  .tab {
+    position: relative;
+    display: inline-flex;
+    flex-direction: column;
+    align-items: center;
+    justify-content: space-between;
+    color: white;
+    display: flex;
+    flex-flow: row;
+    gap: 8px;
+
+    padding: 0;
+    margin: 0;
+    border: none;
+    background: none;
+
+    transition: font-weight 100ms linear;
+
+    & button {
+      cursor: pointer;
+      border: 1px solid var(--gray-500);
+      padding: 3px 5px;
+      border-radius: 5px;
+      line-height: 1;
+
+      background-color: color-mix(in lab, white 80%, transparent 20%);
+      color: var(--background-color);
+
+      &::after {
+        content: attr(data-text);
+        height: 0;
+        visibility: hidden;
+        pointer-events: none;
+        overflow: hidden;
+        font-weight: 700;
+      }
+
+      &:hover {
+        background-color: color-mix(in lab, white 70%, transparent 30%);
+      }
+
+      &:active {
+        background-color: color-mix(in lab, white 60%, transparent 40%);
+      }
+
+      &[selected] {
+        font-weight: 800;
+        color: white;
+        border: 2px solid var(--gray-200);
+        background-color: color-mix(in lab, white 20%, transparent 80%);
+
+        &::before {
+          content: '';
+          z-index: -1;
+          width: 108%;
+          position: absolute;
+          bottom: -4px;
+          left: 50%;
+          transform: translate(-50%, 0);
+
+          border-radius: 5px;
+        }
+      }
+    }
+  }
+}
diff --git a/examples/rag-playground/src/components/app/app.ts b/examples/rag-playground/src/components/app/app.ts
@@ -5,6 +5,7 @@ import { unsafeHTML } from 'lit/directives/unsafe-html.js';
 import '../playground/playground';
 import componentCSS from './app.css?inline';
 import logoIcon from '../../images/icon-logo.svg?raw';
+import { Dataset } from '../playground/playground';
 
 /**
  * App element.
@@ -15,12 +16,15 @@ export class MememoRagPlayground extends LitElement {
   //==========================================================================||
   //                              Class Properties                            ||
   //==========================================================================||
+  @state()
+  curDataset: Dataset;
 
   //==========================================================================||
   //                             Lifecycle Methods                            ||
   //==========================================================================||
   constructor() {
     super();
+    this.curDataset = Dataset.arXiv10k;
   }
 
   /**
@@ -37,6 +41,11 @@ export class MememoRagPlayground extends LitElement {
   //==========================================================================||
   //                              Event Handlers                              ||
   //==========================================================================||
+  tabButtonClicked(dataset: Dataset) {
+    if (dataset !== this.curDataset) {
+      this.curDataset = dataset;
+    }
+  }
 
   //==========================================================================||
   //                             Private Helpers                              ||
@@ -46,6 +55,59 @@ export class MememoRagPlayground extends LitElement {
   //                           Templates and Styles                           ||
   //==========================================================================||
   render() {
+    // Compile the playground
+    let playground = html``;
+
+    switch (this.curDataset) {
+      case Dataset.arXiv10k: {
+        playground = html`<mememo-playground
+          curDataset=${this.curDataset}
+        ></mememo-playground>`;
+        break;
+      }
+
+      case Dataset.arXiv120k: {
+        playground = html`<mememo-playground
+          curDataset=${this.curDataset}
+        ></mememo-playground>`;
+        break;
+      }
+
+      case Dataset.DiffusionDB10k: {
+        playground = html`<mememo-playground
+          curDataset=${this.curDataset}
+        ></mememo-playground>`;
+        break;
+      }
+
+      case Dataset.DiffusionDB100k: {
+        playground = html`<mememo-playground
+          curDataset=${this.curDataset}
+        ></mememo-playground>`;
+        break;
+      }
+
+      case Dataset.DiffusionDB1m: {
+        playground = html`<mememo-playground
+          curDataset=${this.curDataset}
+        ></mememo-playground>`;
+        break;
+      }
+
+      case Dataset.accident3k: {
+        playground = html`<mememo-playground
+          curDataset=${this.curDataset}
+        ></mememo-playground>`;
+        break;
+      }
+
+      default: {
+        playground = html`<mememo-playground
+          curDataset=${Dataset.arXiv10k}
+        ></mememo-playground>`;
+      }
+    }
+
     return html`
       <div class="page">
         <div class="main-app">
@@ -75,8 +137,67 @@ export class MememoRagPlayground extends LitElement {
               </div>
             </div>
 
-            <mememo-playground></mememo-playground>
-            <div class="app-tabs"></div>
+            ${playground}
+
+            <div class="app-tabs">
+              <div class="tab">
+                ML arXiv Abstracts
+                <button
+                  ?selected=${this.curDataset === Dataset.arXiv10k}
+                  @click=${() => this.tabButtonClicked(Dataset.arXiv10k)}
+                >
+                  10k
+                </button>
+                <button
+                  ?selected=${this.curDataset === Dataset.arXiv120k}
+                  @click=${() => this.tabButtonClicked(Dataset.arXiv120k)}
+                >
+                  120k
+                </button>
+              </div>
+
+              <div class="splitter"></div>
+
+              <div class="tab">
+                DiffusionDB Prompts
+                <button
+                  ?selected=${this.curDataset === Dataset.DiffusionDB10k}
+                  @click=${() => this.tabButtonClicked(Dataset.DiffusionDB10k)}
+                >
+                  10k
+                </button>
+                <button
+                  ?selected=${this.curDataset === Dataset.DiffusionDB100k}
+                  @click=${() => this.tabButtonClicked(Dataset.DiffusionDB100k)}
+                >
+                  100k
+                </button>
+                <button
+                  ?selected=${this.curDataset === Dataset.DiffusionDB500k}
+                  @click=${() => this.tabButtonClicked(Dataset.DiffusionDB500k)}
+                >
+                  500k
+                </button>
+                <button
+                  ?selected=${this.curDataset === Dataset.DiffusionDB1m}
+                  @click=${() => this.tabButtonClicked(Dataset.DiffusionDB1m)}
+                >
+                  1M
+                </button>
+              </div>
+
+              <div class="splitter"></div>
+
+              <div class="tab">
+                AI Accident Reports
+                <button
+                  ?selected=${this.curDataset === Dataset.accident3k}
+                  @click=${() => this.tabButtonClicked(Dataset.accident3k)}
+                >
+                  3k
+                </button>
+              </div>
+            </div>
           </div>
 
           <div class="text-right"></div>

diff --git a/examples/rag-playground/src/components/playground/playground.ts b/examples/rag-playground/src/components/playground/playground.ts
@@ -33,16 +33,21 @@ import TextGenLocalWorkerInline from '../../llms/web-llm?worker&inline';
 import EmbeddingWorkerInline from '../../workers/embedding?worker';
 import logoIcon from '../../images/icon-logo.svg?raw';
 
+const STORE_ENDPOINT = 'https://pub-4eccf317e01e4aa3a5caa9991c8b1e2a.r2.dev/';
+
 interface DatasetInfo {
   dataURL: string;
   indexURL?: string;
   datasetName: string;
   datasetNameDisplay: string;
 }
 
-enum Dataset {
+export enum Dataset {
   arXiv10k = 'arxiv-10k',
   arXiv120k = 'arxiv-120k',
+  DiffusionDB10k = 'diffusiondb-10k',
+  DiffusionDB100k = 'diffusiondb-100k',
+  DiffusionDB500k = 'diffusiondb-500k',
   DiffusionDB1m = 'diffusiondb-1m',
   accident3k = 'accident-3k'
 }
@@ -58,29 +63,50 @@ const promptTemplate = promptTemplates as Record<Dataset, string>;
 
 const datasets: Record<Dataset, DatasetInfo> = {
   [Dataset.arXiv10k]: {
-    indexURL: '/data/ml-arxiv-papers-10k-index.json.gzip',
-    dataURL: '/data/ml-arxiv-papers-10k.ndjson.gzip',
+    indexURL: STORE_ENDPOINT + 'ml-arxiv-papers-index-10k.json.gzip',
+    dataURL: STORE_ENDPOINT + 'ml-arxiv-papers-10k.ndjson.gzip',
     datasetName: 'ml-arxiv-papers-10k',
     datasetNameDisplay: 'ML arXiv Abstracts (10k)'
   },
 
   [Dataset.arXiv120k]: {
-    indexURL: '/data/ml-arxiv-papers-120k-index.json.gzip',
-    dataURL: '/data/ml-arxiv-papers-120k.ndjson.gzip',
+    indexURL: STORE_ENDPOINT + 'ml-arxiv-papers-index-120k.json.gzip',
+    dataURL: STORE_ENDPOINT + 'ml-arxiv-papers-120k.ndjson.gzip',
     datasetName: 'ml-arxiv-papers-120k',
     datasetNameDisplay: 'ML arXiv Abstracts (120k)'
   },
 
+  [Dataset.DiffusionDB10k]: {
+    indexURL: STORE_ENDPOINT + 'diffusiondb-prompt-index-10k.json.gzip',
+    dataURL: STORE_ENDPOINT + 'diffusiondb-prompt-10k.ndjson.gzip',
+    datasetName: 'diffusiondb-prompts-10k',
+    datasetNameDisplay: 'DiffusionDB Prompts (10k)'
+  },
+
+  [Dataset.DiffusionDB100k]: {
+    indexURL: STORE_ENDPOINT + 'diffusiondb-prompt-index-100k.json.gzip',
+    dataURL: STORE_ENDPOINT + 'diffusiondb-prompt-100k.ndjson.gzip',
+    datasetName: 'diffusiondb-prompts-100k',
+    datasetNameDisplay: 'DiffusionDB Prompts (100k)'
+  },
+
+  [Dataset.DiffusionDB500k]: {
+    // indexURL: '/data/diffusiondb-prompts-index-1m.json.gzip',
+    dataURL: '/data/diffusiondb-prompt-500k.ndjson.gzip',
+    datasetName: 'diffusiondb-prompts-500k',
+    datasetNameDisplay: 'DiffusionDB Prompts (500k)'
+  },
+
   [Dataset.DiffusionDB1m]: {
-    // indexURL: '/data/diffusiondb-prompts-1m-index.json.gzip',
+    // indexURL: '/data/diffusiondb-prompts-index-1m.json.gzip',
     dataURL: '/data/diffusiondb-prompt-1m.ndjson.gzip',
     datasetName: 'diffusiondb-prompts-1m',
     datasetNameDisplay: 'DiffusionDB Prompts (1M)'
   },
 
   [Dataset.accident3k]: {
-    indexURL: '/data/accident-3k-index.json.gzip',
-    dataURL: '/data/accident-3k.ndjson.gzip',
+    indexURL: STORE_ENDPOINT + 'accident-index-3k.json.gzip',
+    dataURL: STORE_ENDPOINT + 'accident-3k.ndjson.gzip',
     datasetName: 'accidents-3k',
     datasetNameDisplay: 'AI Accidents (3k)'
   }
@@ -212,6 +238,13 @@ export class MememoPlayground extends LitElement {
    */
   willUpdate(changedProperties: PropertyValues<this>) {}
 
+  disconnectedCallback(): void {
+    // Need to terminal workers to avoid aysnc write to indexedDB
+    this.embeddingWorker.terminate();
+    this.textGenLocalWorker.terminate();
+    this.textViewerComponent?.mememoWorker.terminate();
+  }
+
   //==========================================================================||
   //                              Custom Methods                              ||
   //==========================================================================||

diff --git a/examples/rag-playground/src/components/text-viewer/text-viewer.css b/examples/rag-playground/src/components/text-viewer/text-viewer.css
@@ -207,7 +207,7 @@
   .item {
     line-height: 1.25;
     width: 100%;
-    padding: 4px calc(var(--search-inner-padding-h) + var(--box-padding-h));
+    padding: 5px calc(var(--search-inner-padding-h) + var(--box-padding-h));
     border-bottom: 1px solid var(--gray-200);
     cursor: default;
     box-sizing: border-box;
@@ -218,12 +218,14 @@
     }
 
     &[clamp-line] {
-      display: -webkit-box;
-      -webkit-box-orient: vertical;
-      -webkit-line-clamp: 4;
-      overflow: hidden;
-      flex-shrink: 0;
-      -webkit-box-pack: end;
+      .item-content {
+        display: -webkit-box;
+        -webkit-box-orient: vertical;
+        -webkit-line-clamp: 4;
+        overflow: hidden;
+        flex-shrink: 0;
+        -webkit-box-pack: end;
+      }
     }
 
     &:hover {