Allow users to change ef-search

Signed-off-by: Jay Wang <jay@zijie.wang>

Author: xiaohk

examples/rag-playground/public/data/ml-arxiv-papers-1000-index.json

@@ -59,6 +59,13 @@ export class MememoRagPlayground extends LitElement {
     let playground = html``;
 
     switch (this.curDataset) {
+      case Dataset.arXiv1k: {
+        playground = html`<mememo-playground
+          curDataset=${this.curDataset}
+        ></mememo-playground>`;
+        break;
+      }
+
       case Dataset.arXiv10k: {
         playground = html`<mememo-playground
           curDataset=${this.curDataset}
@@ -142,6 +149,12 @@ export class MememoRagPlayground extends LitElement {
             <div class="app-tabs">
               <div class="tab">
                 ML arXiv Abstracts
+                <button
+                  ?selected=${this.curDataset === Dataset.arXiv1k}
+                  @click=${() => this.tabButtonClicked(Dataset.arXiv1k)}
+                >
+                  1k
+                </button>
                 <button
                   ?selected=${this.curDataset === Dataset.arXiv10k}
                   @click=${() => this.tabButtonClicked(Dataset.arXiv10k)}

examples/rag-playground/public/data/ml-arxiv-papers-1000.ndjson

@@ -408,6 +408,7 @@
   text-align: center;
   position: relative;
   top: -10px;
+  gap: 10px;
 }
 
 .search-top-info {
@@ -448,3 +449,67 @@
 input {
   all: unset;
 }
+
+.svg-icon {
+  display: flex;
+  justify-content: center;
+  align-items: center;
+  width: 1em;
+  height: 1em;
+
+  color: currentColor;
+  transition: transform 80ms linear;
+  transform-origin: center;
+
+  & svg {
+    fill: currentColor;
+    width: 100%;
+    height: 100%;
+  }
+}
+
+.button-group {
+  display: flex;
+  flex-flow: row;
+  align-items: center;
+  gap: 8px;
+}
+
+button {
+  all: unset;
+
+  display: flex;
+  line-height: 1;
+  padding: 4px 6px;
+  border-radius: 5px;
+  white-space: nowrap;
+
+  cursor: pointer;
+  user-select: none;
+  -webkit-user-select: none;
+
+  background-color: color-mix(in lab, var(--gray-200), white 20%);
+  color: var(--gray-800);
+  display: flex;
+  flex-flow: row;
+  align-items: center;
+  font-size: var(--header-secondary-size);
+  height: 1em;
+
+  &:hover {
+    background-color: color-mix(in lab, var(--gray-300), white 30%);
+  }
+
+  &:active {
+    background-color: color-mix(in lab, var(--gray-300), white 20%);
+  }
+
+  .svg-icon {
+    position: relative;
+    top: 1px;
+    margin-right: 3px;
+    color: var(--gray-700);
+    width: 12px;
+    height: 12px;
+  }
+}

examples/rag-playground/public/data/ml-arxiv-papers-1000.ndjson.gzip

@@ -31,9 +31,11 @@ import '../text-viewer/text-viewer';
 import componentCSS from './playground.css?inline';
 import TextGenLocalWorkerInline from '../../llms/web-llm?worker&inline';
 import EmbeddingWorkerInline from '../../workers/embedding?worker';
-import logoIcon from '../../images/icon-logo.svg?raw';
+import gearIcon from '../../images/icon-gear.svg?raw';
 
-const STORE_ENDPOINT = 'https://pub-4eccf317e01e4aa3a5caa9991c8b1e2a.r2.dev/';
+const STORE_ENDPOINT = import.meta.env.DEV
+  ? '/data/'
+  : 'https://pub-4eccf317e01e4aa3a5caa9991c8b1e2a.r2.dev/';
 
 interface DatasetInfo {
   dataURL: string;
@@ -43,6 +45,7 @@ interface DatasetInfo {
 }
 
 export enum Dataset {
+  arXiv1k = 'arxiv-1k',
   arXiv10k = 'arxiv-10k',
   arXiv120k = 'arxiv-120k',
   DiffusionDB10k = 'diffusiondb-10k',
@@ -62,6 +65,13 @@ enum Arrow {
 const promptTemplate = promptTemplates as Record<Dataset, string>;
 
 const datasets: Record<Dataset, DatasetInfo> = {
+  [Dataset.arXiv1k]: {
+    indexURL: STORE_ENDPOINT + 'ml-arxiv-papers-index-1k.json.gzip',
+    dataURL: STORE_ENDPOINT + 'ml-arxiv-papers-1k.ndjson.gzip',
+    datasetName: 'ml-arxiv-papers-1k',
+    datasetNameDisplay: 'ML arXiv Abstracts (1k)'
+  },
+
   [Dataset.arXiv10k]: {
     indexURL: STORE_ENDPOINT + 'ml-arxiv-papers-index-10k.json.gzip',
     dataURL: STORE_ENDPOINT + 'ml-arxiv-papers-10k.ndjson.gzip',
@@ -147,7 +157,7 @@ export class MememoPlayground extends LitElement {
   topK = 10;
 
   @state()
-  maxDistance = 0.25;
+  efSearch = 100;
 
   @query('mememo-text-viewer')
   textViewerComponent: MememoTextViewer | undefined | null;
@@ -280,7 +290,12 @@ export class MememoPlayground extends LitElement {
     // Loading the arrow
     this.arrowElements[Arrow.Search]?.classList.add('running');
 
-    this.textViewerComponent.semanticSearch(embedding, this.topK, 0.5);
+    this.textViewerComponent.semanticSearch(
+      embedding,
+      this.topK,
+      0.5,
+      this.efSearch
+    );
   }
 
   //==========================================================================||
@@ -346,17 +361,17 @@ export class MememoPlayground extends LitElement {
    * @param e Input event
    * @param parameter Type of the parameter
    */
-  parameterInputChanged(e: InputEvent, parameter: 'distance' | 'top-k') {
+  parameterInputChanged(e: InputEvent, parameter: 'efSearch' | 'top-k') {
     const element = e.currentTarget as HTMLInputElement;
 
-    if (parameter === 'distance') {
-      const value = parseFloat(element.value);
-      if (value > 0 && value < 1) {
-        this.maxDistance = value;
+    if (parameter === 'efSearch') {
+      const value = parseInt(element.value);
+      if (value > 0) {
+        this.efSearch = value;
       } else {
-        this.maxDistance = 0.25;
+        this.efSearch = 100;
       }
-      element.value = String(this.maxDistance);
+      element.value = String(this.efSearch);
     }
 
     if (parameter === 'top-k') {
@@ -574,13 +589,13 @@ export class MememoPlayground extends LitElement {
           <div class="search-box">
             <div class="search-top-info">
               <span class="row"
-                >distance &lt;
+                >ef-search =
                 <input
-                  id="input-distance"
+                  id="input-efSearch"
                   type="text"
-                  value="0.25"
+                  value="100"
                   @change=${(e: InputEvent) =>
-                    this.parameterInputChanged(e, 'distance')}
+                    this.parameterInputChanged(e, 'efSearch')}
               /></span>
               <span class="row"
                 >top-k =
@@ -623,6 +638,12 @@ export class MememoPlayground extends LitElement {
         <div class="container container-model">
           <div class="model-box">
             <div class="header">GPT 3.5</div>
+            <div class="button-group">
+              <button>
+                <span class="svg-icon">${unsafeHTML(gearIcon)}</span>
+                change
+              </button>
+            </div>
           </div>
         </div>
 

examples/rag-playground/src/components/app/app.ts

@@ -147,14 +147,20 @@ export class MememoTextViewer extends LitElement {
    * @param topK Top k relevant documents to retrieve
    * @param maxDistance Distance threshold for relevance
    */
-  semanticSearch(embedding: number[], topK: number, maxDistance: number) {
+  semanticSearch(
+    embedding: number[],
+    topK: number,
+    maxDistance: number,
+    efSearch: number
+  ) {
     const message: MememoWorkerMessage = {
       command: 'startSemanticSearch',
       payload: {
         embedding,
         requestID: this.semanticSearchRequestID,
         topK,
-        maxDistance
+        maxDistance,
+        efSearch
       }
     };
     this.mememoWorker.postMessage(message);

examples/rag-playground/src/components/playground/playground.css

@@ -74,6 +74,7 @@ export type MememoWorkerMessage =
         requestID: number;
         topK: number;
         maxDistance: number;
+        efSearch: number;
       };
     }
   | {
@@ -83,6 +84,7 @@ export type MememoWorkerMessage =
         requestID: number;
         topK: number;
         maxDistance: number;
+        efSearch: number;
         documents: string[];
         documentDistances: number[];
       };
@@ -165,8 +167,9 @@ self.onmessage = (e: MessageEvent<MememoWorkerMessage>) => {
     }
 
     case 'startSemanticSearch': {
-      const { requestID, embedding, topK, maxDistance } = e.data.payload;
-      semanticSearch(embedding, topK, maxDistance, requestID).then(
+      const { requestID, embedding, topK, efSearch, maxDistance } =
+        e.data.payload;
+      semanticSearch(embedding, topK, maxDistance, efSearch, requestID).then(
         () => {},
         () => {}
       );
@@ -381,6 +384,7 @@ const semanticSearch = async (
   embedding: number[],
   topK: number,
   maxDistance: number,
+  efSearch: number,
   requestID: number
 ) => {
   if (documentDBPromise === null) {
@@ -389,7 +393,7 @@ const semanticSearch = async (
   const documentDB = await documentDBPromise;
   await finishedLoading;
 
-  const { keys, distances } = await hnswIndex.query(embedding, topK);
+  const { keys, distances } = await hnswIndex.query(embedding, topK, efSearch);
 
   // Batched query documents from indexedDB
   const results = await documentDB.bulkGet(keys);
@@ -414,6 +418,7 @@ const semanticSearch = async (
       topK,
       requestID,
       maxDistance,
+      efSearch,
       documents,
       documentDistances
     }