working

package.json

@@ -15,8 +15,6 @@
         "@types/react": "18.2.6",
         "@types/react-dom": "18.2.4",
         "@xenova/transformers": "^2.1.1",
-        "chromadb": "^1.5.2",
-        "cohere-ai": "^5.1.0",
         "dexie": "^3.2.4",
         "eslint": "8.40.0",
         "eslint-config-next": "13.4.2",

src/components/ChatWindow.jsx

@@ -5,11 +5,8 @@ import MessageList from './MessageList';
 import {FileLoader} from './FileLoader';
 import Loader from "./Loader";
 import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
-import { TransformersEmbeddingFunction } from '../embed/hf';
-import { ChromaClient } from "chromadb";
-
-const client = new ChromaClient();
-const embedder = new TransformersEmbeddingFunction({});
+import { XenovaTransformersEmbeddings } from '../embed/hf';
+import { MemoryVectorStore } from "langchain/vectorstores/memory";
 
 function ChatWindow({
   stopStrings,
@@ -34,37 +31,31 @@ function ChatWindow({
       console.log('found file text splitting into chunks')
       const textSplitter = new RecursiveCharacterTextSplitter({ chunkSize: 1000 });
       const docs = await textSplitter.createDocuments([fileText]);
-      console.log(`split docs: ${docs}`);
-      const collection = await client.createCollection({name: "docs", embeddingFunction: embedder })
-      console.log(`collection: ${collection}`);
       let queryResult;
+      let qaPrompt;
+      console.log(docs);
       try {
-        await collection.add({
-          ids: [...docs.map((v, k) => k)],
-          metadatas: [...docs.map(doc => doc.metadata)],
-          documents: [...docs.map(doc => doc.pageContent)],
-        });
-        const queryResult = await collection.query({
-          nResults: 2, 
-          queryTexts: [userPrompt]
-        });
-        console.log(queryResult);
-      } catch (err) {
-        console.log(err);
-      }
-      
-
-      const qaPrompt = 
+        const vectorStore = await MemoryVectorStore.fromTexts(
+          [...docs.map(doc => doc.pageContent)],
+          [...docs.map((v, k) => k)],
+          new XenovaTransformersEmbeddings()
+        )
+        let queryResult = await vectorStore.similaritySearch(userInput, 1);
+        console.log("queryResult", queryResult);
+        qaPrompt = 
       `You are an AI assistant providing helpful advice. You are given the following extracted parts of a long document and a question. Provide a conversational answer based on the context provided.
         You should only provide hyperlinks that reference the context below. Do NOT make up hyperlinks.
         If you can't find the answer in the context below, just say "Hmm, I'm not sure." Don't try to make up an answer.
         If the question is not related to the context, politely respond that you are tuned to only answer questions that are related to the context.
         Question: ${userInput}
         =========
-        ${queryResult}
+        ${queryResult[0].pageContent}
         =========
         Answer:
         `
+      } catch (err) {
+        console.log(err);
+      }      
       send(qaPrompt, maxTokens, stopStrings);
     } else {
       send(userInput, maxTokens, stopStrings);

src/embed/hf.ts

@@ -1,62 +1,42 @@
-import { IEmbeddingFunction } from "chromadb/src/embeddings/IEmbeddingFunction";
-
-// Dynamically import module
-let TransformersApi: Promise<any>;
-
-export class TransformersEmbeddingFunction implements IEmbeddingFunction {
-  private pipelinePromise: Promise<any> | null;
-
-  /**
-   * TransformersEmbeddingFunction constructor.
-   * @param options The configuration options.
-   * @param options.model The model to use to calculate embeddings. Defaults to 'Xenova/all-MiniLM-L6-v2', which is an ONNX port of `sentence-transformers/all-MiniLM-L6-v2`.
-   * @param options.revision The specific model version to use (can be a branch, tag name, or commit id). Defaults to 'main'.
-   * @param options.quantized Whether to load the 8-bit quantized version of the model. Defaults to `false`.
-   * @param options.progress_callback If specified, this function will be called during model construction, to provide the user with progress updates.
-   */
-  constructor({
-    model = "Xenova/all-MiniLM-L6-v2",
-    revision = "main",
-    quantized = false,
-    progress_callback = null,
-  }: {
-    model?: string;
-    revision?: string;
-    quantized?: boolean;
-    progress_callback?: Function | null;
-  } = {}) {
-    try {
-      // Since Transformers.js is an ESM package, we use the dynamic `import` syntax instead of `require`.
-      // Also, since we use `"module": "commonjs"` in tsconfig.json, we use the following workaround to ensure
-      // the dynamic import is not transpiled to a `require` statement.
-      // For more information, see https://github.com/microsoft/TypeScript/issues/43329#issuecomment-1008361973
-      TransformersApi = Function('return import("@xenova/transformers")')();
-    } catch (e) {
-      throw new Error(
-        "Please install the @xenova/transformers package to use the TransformersEmbeddingFunction, `npm install -S @xenova/transformers`."
-      );
+import { pipeline } from "@xenova/transformers";
+import { Embeddings, EmbeddingsParams } from "langchain/embeddings/base";
+
+export interface XenovaTransformersEmbeddingsParams extends EmbeddingsParams {
+  model?: string;
+}
+
+export class XenovaTransformersEmbeddings
+  extends Embeddings
+  implements XenovaTransformersEmbeddingsParams
+{
+  model: string;
+
+  client: any;
+
+  constructor(fields?: XenovaTransformersEmbeddingsParams) {
+    super(fields ?? {});
+    this.model = fields?.model ?? "Xenova/all-MiniLM-L6-v2";
+  }
+
+  async _embed(texts: string[]): Promise<number[][]> {
+    if (!this.client) {
+      this.client = await pipeline("embeddings", this.model);
     }
 
-    // Store a promise that resolves to the pipeline
-    this.pipelinePromise = new Promise(async (resolve, reject) => {
-      try {
-        const { pipeline } = await TransformersApi;
-        resolve(
-          await pipeline("feature-extraction", model, {
-            quantized,
-            revision,
-            progress_callback,
-          })
-        );
-      } catch (e) {
-        reject(e);
-      }
+    return this.caller.call(async () => {
+      return await Promise.all(
+        texts.map(async (t) => (await this.client(t, {
+          pooling: "mean", normalize: true 
+        })).data)
+      );
     });
   }
 
-  public async generate(texts: string[]): Promise<number[][]> {
-    let pipe = await this.pipelinePromise;
-    let output = await pipe(texts, { pooling: "mean", normalize: true });
-    return output.tolist();
+  embedQuery(document: string): Promise<number[]> {
+    return this._embed([document]).then((embeddings) => embeddings[0]);
+  }
+
+  embedDocuments(documents: string[]): Promise<number[][]> {
+    return this._embed(documents);
   }
 }