✨ models: Add some options to transcribe & add support for array embe…

…ddings (#130) * ✨ models: Add some options to transcribe & add support for array embeddings * ⚠️ models: Fix lint issue * Update lib/transcribe.ts Co-authored-by: Victor Forissier <34337148+victorforissier@users.noreply.github.com> --------- Co-authored-by: Victor Forissier <34337148+victorforissier@users.noreply.github.com>
polyfact · Feb 5, 2024 · 49b2294 · 49b2294
1 parent 58530fc
commit 49b2294
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 7 deletions.
diff --git a/lib/embeddings.ts b/lib/embeddings.ts
@@ -34,7 +34,7 @@ async function createEmbeddings(
 
 async function updateEmbeddings(
     id: string,
-    input: string,
+    input: string | string[],
     maxToken = 0,
     clientOptions: InputClientOptions = {},
 ): Promise<{ success: boolean }> {
@@ -154,7 +154,10 @@ class Embeddings {
         }
     }
 
-    async add(input: string, { maxToken = 0 }: EmbeddingsAddOptions = {}): Promise<void> {
+    async add(
+        input: string | string[],
+        { maxToken = 0 }: EmbeddingsAddOptions = {},
+    ): Promise<void> {
         const id = await this.memoryId;
         await updateEmbeddings(id, input, maxToken, await this.clientOptions);
     }
@@ -173,13 +176,17 @@ export { createEmbeddings, updateEmbeddings, getAllEmbeddings, Embeddings };
 
 export type EmbeddingsClient = {
     createMemory: (isPublic: true) => Promise<{ id: string }>;
-    updateMemory: (id: string, input: string, maxToken?: number) => Promise<{ success: boolean }>;
+    updateMemory: (
+        id: string,
+        input: string | string[],
+        maxToken?: number,
+    ) => Promise<{ success: boolean }>;
     getAllMemories: () => Promise<{ ids: string[] }>;
     Memory: () => Embeddings;
     createEmbeddings: (isPublic: true) => Promise<{ id: string }>;
     updateEmbeddings: (
         id: string,
-        input: string,
+        input: string | string[],
         maxToken?: number,
     ) => Promise<{ success: boolean }>;
     getAllEmbeddings: () => Promise<{ ids: string[] }>;
@@ -189,13 +196,13 @@ export type EmbeddingsClient = {
 export default function client(clientOptions: InputClientOptions = {}): EmbeddingsClient {
     return {
         createMemory: () => createEmbeddings(clientOptions),
-        updateMemory: (id: string, input: string, maxToken?: number) =>
+        updateMemory: (id: string, input: string | string[], maxToken?: number) =>
             updateEmbeddings(id, input, maxToken, clientOptions),
         getAllMemories: () => getAllEmbeddings(clientOptions),
         Memory: (embeddingsOptions?: EmbeddingsOptions) =>
             new Embeddings(embeddingsOptions, clientOptions),
         createEmbeddings: () => createEmbeddings(clientOptions),
-        updateEmbeddings: (id: string, input: string, maxToken?: number) =>
+        updateEmbeddings: (id: string, input: string | string[], maxToken?: number) =>
             updateEmbeddings(id, input, maxToken, clientOptions),
         getAllEmbeddings: () => getAllEmbeddings(clientOptions),
         Embeddings: (embeddingsOptions?: EmbeddingsOptions) =>

diff --git a/lib/transcribe.ts b/lib/transcribe.ts
@@ -21,6 +21,12 @@ const ResultType = t.intersection([
                 speaker_confidence: t.number,
             }),
         ),
+        dialogue: t.array(
+            t.type({
+                speaker: t.number,
+                text: t.string,
+            }),
+        ),
     }),
 ]);
 
@@ -31,6 +37,8 @@ function randomString() {
 
 type TranscriptionOptions = {
     provider?: string;
+    language?: string;
+    outputFormat?: string;
 };
 
 type Word = {
@@ -43,6 +51,11 @@ type Word = {
     speakerConfidence: number;
 };
 
+type DialogueElement = {
+    speaker: number;
+    text: string;
+};
+
 export class Transcription implements Promise<string> {
     [Symbol.toStringTag] = "Transcription";
 
@@ -86,6 +99,19 @@ export class Transcription implements Promise<string> {
             })) ?? []
         );
     }
+
+    async dialogue(): Promise<DialogueElement[]> {
+        const { dialogue } = await this.promise;
+
+        if (!dialogue || dialogue.length === 0) {
+            console.warn(
+                "No dialogue found in transcription even though text has been transcribed. \nThis could indicate the provider you are using does not support word-level transcriptions. If you did not explicitely defined a provider when calling the transcribe function, the default provider is openai and does not support word level transcription.\nYou can switch to a provider like deepgram if you need word-level transcription. For more information check the Polyfire documentation at https://docs.polyfire.com.",
+            );
+            return [];
+        }
+
+        return dialogue;
+    }
 }
 
 export function transcribe(
@@ -110,7 +136,12 @@ export function transcribe(
 
             const res = await axios.post(
                 `${endpoint}/transcribe`,
-                { ...options, file_path: fileName },
+                {
+                    provider: options.provider,
+                    language: options.language,
+                    output_format: options.outputFormat,
+                    file_path: fileName,
+                },
                 {
                     headers: {
                         "X-Access-Token": token,