diff --git a/.gitignore b/.gitignore index 4076ee1..941af50 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ prompt.md todo .toak-ignore +/google_gemma-3-1b-it-Q6_K.llamafile diff --git a/bun.lockb b/bun.lockb index 45b6afc..11d9545 100755 Binary files a/bun.lockb and b/bun.lockb differ diff --git a/package.json b/package.json index 403f766..7368727 100644 --- a/package.json +++ b/package.json @@ -45,7 +45,9 @@ "lint:fix": "eslint src/ --fix", "format": "prettier --write \"**/*.{js,jsx,ts,tsx,json,md,yml,yaml}\"", "fix": "bun format && bun lint:fix", - "release": "bunx release-it" + "release": "bunx release-it", + "run:inference": "./google_gemma-3-1b-it-Q6_K.llamafile --server --nobrowser", + "download:inference": "./download_optimization_model.sh" }, "dependencies": { "glob": "^11.0.1", diff --git a/scripts/download_optimization_model.sh b/scripts/download_optimization_model.sh new file mode 100644 index 0000000..8274553 --- /dev/null +++ b/scripts/download_optimization_model.sh @@ -0,0 +1,4 @@ +MODEL_URL=https://huggingface.co/Mozilla/gemma-3-1b-it-llamafile/resolve/main/google_gemma-3-1b-it-Q6_K.llamafile?download=true + +wget "${MODEL_URL}" + diff --git a/scripts/optimize-ignore.ts b/scripts/optimize-ignore.ts new file mode 100755 index 0000000..fa49074 --- /dev/null +++ b/scripts/optimize-ignore.ts @@ -0,0 +1,30 @@ +export async function optimizeToakIgnore(content: string) { + const inferenceProcess = Bun.spawn(['bun', 'run:inference']); + + await new Promise(resolve => setTimeout(resolve, 5000)); + + const prompt = `You are a helpful assistant. +## Context +~~~ +${content} +~~~ +Respond with a list of files that should be added to the .toak-ignore file to reduce noise in the context. No extra text or explanations.`; + + async function run() { + try { + const response = await fetch('http://127.0.0.1:8080/completion', { + method: 'POST', + body: JSON.stringify({ + prompt, + n_predict: 512, + }), + }); + const data = await response.json(); + console.log(data.content); + } catch (error) { + console.error('Error:', error); + } + } + await run(); + inferenceProcess.kill(); +} \ No newline at end of file diff --git a/src/cli.ts b/src/cli.ts index 8d519b4..e9c63fc 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -3,10 +3,22 @@ import type { PresetPrompt } from './prompts'; console.log('RUNNING TOKENIZER'); import { MarkdownGenerator, type MarkdownGeneratorOptions } from './MarkdownGenerator'; +import { optimizeToakIgnore } from '../scripts/optimize-ignore.ts'; + +let output = ''; +process.stdout.write = (write => { + return (str: string | Uint8Array, ...args: any) => { + output += str.toString(); + return write.apply(process.stdout, [str, ...args]); + }; +})(process.stdout.write); + +let optimizeIgnore = false; const args = process.argv.slice(2); const options: { prompt?: PresetPrompt; } & MarkdownGeneratorOptions = { + }; type ValidArg = keyof MarkdownGeneratorOptions; @@ -18,6 +30,10 @@ for (let i = 0; i < args.length; i++) { options["todoPrompt"] = args[i + 1] i++; } + if (args[i] === '--optimize-ignore') { + optimizeIgnore = true; + i++; + } const arg = args[i].replace(/^--/, ''); if (arg as any satisfies ValidArg) { // @ts-ignore - arg can't be used to index options @@ -30,12 +46,19 @@ for (let i = 0; i < args.length; i++) { } const generator = new MarkdownGenerator(options); + + + + generator .createMarkdownDocument() - .then((result: { success: boolean }) => { + .then(async (result: { success: boolean }) => { if (!result.success) { process.exit(1); } + if (optimizeIgnore) { + await optimizeToakIgnore(output); + } }) .catch((error: any) => { console.error('Error:', error);