feat(website): add Embedding dev tool

aws-samples · Nov 7, 2023 · 30485b7 · 30485b7
1 parent db79a02
commit 30485b7
Show file tree

Hide file tree

Showing 7 changed files with 124 additions and 39 deletions.
diff --git a/demo/corpus/embeddings/poetry.lock b/demo/corpus/embeddings/poetry.lock
diff --git a/demo/website/.projen/deps.json b/demo/website/.projen/deps.json
diff --git a/demo/website/.projen/tasks.json b/demo/website/.projen/tasks.json
diff --git a/demo/website/package.json b/demo/website/package.json
diff --git a/demo/website/src/pages/Embeddings/index.tsx b/demo/website/src/pages/Embeddings/index.tsx
@@ -14,20 +14,45 @@ import {
 } from '@cloudscape-design/components';
 import ContentLayout from '@cloudscape-design/components/content-layout';
 import Header from '@cloudscape-design/components/header';
-import { useEmbedQuery } from 'api-typescript-react-query-hooks';
-import { useCallback, useState } from 'react';
+import { faker } from '@faker-js/faker';
+import { useEmbedDocuments } from 'api-typescript-react-query-hooks';
+import { useState } from 'react';
+import { useImmer } from 'use-immer';
 
 export const EmbeddingsTool: React.FC = () => {
-  const [query, setQuery] = useState<string>('');
-  const embedQuery = useEmbedQuery();
+  const [sentences, updateSentences] = useImmer<string[]>(['text']);
+  const [duration, setDuration] = useState<number>();
+  const caller = useEmbedDocuments();
+  const sampleCount = 20;
 
-  const onSubmit = useCallback(() => {
-    embedQuery.mutate({
-      embedQueryRequestContent: {
-        text: query,
+  const onSubmit = async () => {
+    const _start = Date.now();
+    setDuration(undefined);
+    await caller.mutateAsync({
+      embedDocumentsRequestContent: {
+        texts: sentences,
       },
     });
-  }, [embedQuery, query]);
+    setDuration((Date.now() - _start) / 1000);
+  };
+
+  const addSentence = () => {
+    updateSentences((draft) => {
+      draft.push('text');
+    });
+  };
+  const removeSentence = () => {
+    updateSentences((draft) => {
+      if (sentences.length > 1) {
+        draft.pop();
+      }
+    });
+  };
+  const loadSampleData = () => {
+    updateSentences(() => {
+      return Array.from(Array(sampleCount)).map(() => faker.lorem.sentence(50));
+    });
+  };
 
   return (
     <ContentLayout
@@ -42,35 +67,75 @@ export const EmbeddingsTool: React.FC = () => {
           <Form
             actions={
               <SpaceBetween direction="horizontal" size="s">
-                <Button onClick={onSubmit} loading={embedQuery.isLoading}>
+                <Button onClick={removeSentence} variant="inline-icon" iconName="undo" disabled={caller.isLoading}>
+                  -
+                </Button>
+                <Button onClick={addSentence} variant="inline-icon" iconName="add-plus" disabled={caller.isLoading}>
+                  +
+                </Button>
+                <Button onClick={loadSampleData} variant="link" disabled={caller.isLoading}>
+                  Sample
+                </Button>
+                <Button onClick={onSubmit} variant="primary" loading={caller.isLoading}>
                   Submit
                 </Button>
               </SpaceBetween>
             }
           >
-            <FormField label="Input Text" description="Enter the text to embed">
-              <Textarea value={query} onChange={({ detail }) => setQuery(detail.value)} />
+            <FormField label="Input Text" description="Enter the text to embed" stretch>
+              <div
+                style={{
+                  maxHeight: '25vh',
+                  overflow: 'auto',
+                  paddingRight: 60,
+                }}
+              >
+                <SpaceBetween direction="vertical" size="m">
+                  {sentences.map((s, i) => (
+                    <Textarea
+                      key={i}
+                      value={s}
+                      onChange={({ detail }) =>
+                        updateSentences((draft) => {
+                          draft[i] = detail.value;
+                        })
+                      }
+                    />
+                  ))}
+                </SpaceBetween>
+              </div>
             </FormField>
           </Form>
         </Container>
 
-        <Box variant="div">
-          <Header counter={`(${embedQuery.data?.embedding.length || 0})`} info={embedQuery.data?.model}>
-            Results
-          </Header>
-          <br />
-          {embedQuery.isLoading && <Spinner />}
+        <SpaceBetween direction="vertical" size="l">
+          {caller.isLoading && <Spinner />}
 
-          {embedQuery.error && (
+          {caller.error && (
             <Alert type="error" header="Error">
-              <code>{embedQuery.error.message}</code>
+              <code>{caller.error.message}</code>
             </Alert>
           )}
 
-          <TextContent>
-            <div style={{ whiteSpace: 'pre-wrap' }}>{JSON.stringify(embedQuery.data?.embedding, null, 0)}</div>
-          </TextContent>
-        </Box>
+          <div>{caller.data && `${sentences.length} documents in ${duration} seconds`}</div>
+
+          <div style={{ maxHeight: '25vh', overflow: 'auto', width: '100%' }}>
+            <SpaceBetween direction="vertical" size="m">
+              {caller.data?.embeddings.map((embedding, i) => (
+                <Box key={i} variant="div">
+                  <Header counter={`(${embedding.length || 0})`} info={caller.data?.model}>
+                    Results
+                  </Header>
+                  <br />
+
+                  <TextContent>
+                    <div style={{ whiteSpace: 'pre-wrap' }}>{JSON.stringify(embedding, null, 0)}</div>
+                  </TextContent>
+                </Box>
+              ))}
+            </SpaceBetween>
+          </div>
+        </SpaceBetween>
       </SpaceBetween>
     </ContentLayout>
   );