feat: ✨ add line tracking and summarization script

microsoft · Oct 23, 2024 · 2ff37c9 · 2ff37c9
1 parent 7c7e185
commit 2ff37c9
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 6 deletions.
diff --git a/packages/core/src/encoders.ts b/packages/core/src/encoders.ts
@@ -73,13 +73,17 @@ export async function chunk(
         keepSeparators: true,
     })
     const chunksRaw = ts.split(content)
-    const chunks = chunksRaw.map(({ text, startPos }) => {
+    const chunks = chunksRaw.map(({ text, startPos, endPos }) => {
+        const lineStart = indexToLineNumber(content, startPos)
+        const lineEnd = indexToLineNumber(content, endPos)
         if (lineNumbers) {
-            const startLine = indexToLineNumber(content, startPos)
-            text = addLineNumbers(text, { startLine })
+            text = addLineNumbers(text, { startLine: lineStart })
         }
         return {
-            text,
+            content: text,
+            filename,
+            lineStart,
+            lineEnd,
         } satisfies TextChunk
     })
     return chunks

diff --git a/packages/core/src/types/prompt_template.d.ts b/packages/core/src/types/prompt_template.d.ts
@@ -1095,8 +1095,9 @@ interface CSVParseOptions {
     headers?: string[]
 }
 
-interface TextChunk {
-    text: string
+interface TextChunk extends WorkspaceFile {
+    lineStart: number
+    lineEnd: number
 }
 
 interface TextChunkerConfig extends LineNumberingOptions {

diff --git a/packages/sample/genaisrc/chunk.genai.mjs b/packages/sample/genaisrc/chunk.genai.mjs
@@ -0,0 +1,24 @@
+script({
+    files: "src/rag/loremipsum.pdf",
+})
+
+const chunks = await tokenizers.chunk(env.files[0], {
+    chunkSize: 256,
+    chunkOverlap: 42,
+    lineNumbers: true,
+})
+
+let summary = ""
+for (const chunk of chunks) {
+    const { text } = await runPrompt(
+        (ctx) => {
+            ctx.def("CHUNK", chunk)
+            ctx.def("SUMMARY_SO_FAR", summary, { ignoreEmpty: true })
+            ctx.$`Summarize CHUNK. Use SUMMARY_SO_FAR as a starting point (but do not repeat it).`
+        },
+        { model: "small", system: ["system"] }
+    )
+    summary = text
+}
+
+console.log(summary)