@imaduddin/fastbrain 0.0.0-stage → 1.0.0

package/src/Softmax/index.js

@@ -0,0 +1,225 @@
+import { log } from 'console';
+import { OperationManager } from '../OperationManager.js'
+
+export class Softmax {
+    device;
+    vocabSize;
+    forwardPipline;
+    embeddingSize;
+    gpu;
+    paramsForward;
+    oldParam;
+    bp;
+    totalLoss;
+    totalAcc;
+    totalCalculate;
+    constructor(device, embeddingSize, vocabSize) {
+        this.vocabSize = vocabSize;
+        this.device = device;
+        this.embeddingSize = embeddingSize;
+        this.gpu = new OperationManager(device);
+        this.forwardPipline = this.gpu.getShaderPipline("./src/shaders/SF.wgsl")
+        this.backwardPipline = this.gpu.getShaderPipline("./src/shaders/SB.wgsl")
+        this.lossPipline = this.gpu.getShaderPipline("./src/shaders/CalculateLoss.wgsl")
+        this.accPipline = this.gpu.getShaderPipline("./src/shaders/CalculateAcc.wgsl")
+        this.battention = this.gpu.getShaderPipline("./src/shaders/SoftmaxBackward.wgsl")
+        this.reduceFn = this.gpu.getShaderPipline("./src/shaders/ReductionFunction.wgsl")
+        this.totalCalculate = this.gpu.getShaderPipline('./src/shaders/calculateTotal.wgsl')
+    }
+    async forward(logits, encoderCPU) {
+        const encoder = encoderCPU ? encoderCPU : this.device.createCommandEncoder()
+        const tokenCount = logits.size / (this.vocabSize * 4)
+        if (!this.paramsForward) {
+            this.paramsForward = await this.gpu.createUniFormBuffer(4 * 3)
+        }
+        if (!this.oldParam) {
+            this.oldParam = new Uint32Array([this.vocabSize, 0, 0])
+        }
+
+        if (this.oldParam[1] !== tokenCount) { await this.gpu.WriteBuffer(this.paramsForward, new Uint32Array([this.vocabSize, tokenCount, 0])); this.oldParam = new Uint32Array([this.vocabSize, tokenCount, 0]); }
+        if (!this.output) this.output = await this.gpu.ManageBuffer(logits.size);
+        if (this.output.size !== 4 * tokenCount * this.vocabSize || this.output.isDestroyed()) {
+            this.output.destroy()
+            this.output = await this.gpu.ManageBuffer(logits.size)
+        }
+        await this.output.count()
+        const bind = await this.device.createBindGroup({
+            layout: await this.forwardPipline.getBindGroupLayout(0),
+            entries: [
+                {
+                    binding: 0,
+                    resource: {
+                        buffer: logits.mainBuffer
+                    }
+                },
+                {
+                    binding: 1,
+                    resource: {
+                        buffer: this.output.mainBuffer
+                    }
+                }, {
+                    binding: 2,
+                    resource: {
+                        buffer: this.paramsForward
+                    }
+                }
+            ]
+        })
+        await this.gpu.RunPipline(encoder, this.forwardPipline, bind, tokenCount)
+        if (!encoderCPU) {
+            await this.device.queue.submit([encoder.finish()])
+        }
+        return this.output
+    }
+    async backward(probs, target, encoderMain) {
+        const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder()
+        const tokenCount = probs.size / (this.vocabSize * 4)
+        if (!this.paramsForward) {
+            this.paramsForward = await this.gpu.createUniFormBuffer(4 * 3)
+        }
+        if (!this.oldParam) {
+            this.oldParam = new Uint32Array([this.vocabSize, 0, 0])
+        }
+        if (this.oldParam[1] !== tokenCount) {
+            await this.gpu.WriteBuffer(this.paramsForward, new Uint32Array([this.vocabSize, tokenCount, 0]));
+            this.oldParam = new Uint32Array([this.vocabSize, tokenCount, 0]);
+        }
+        if (!this.targetBuffer) this.targetBuffer = await this.gpu.createBuffer(target.length * Float32Array.BYTES_PER_ELEMENT);
+        if (this.targetBuffer.size !== target.length * 4) {
+            this.targetBuffer.destroy()
+            this.targetBuffer = await this.gpu.createBuffer(target.length * Float32Array.BYTES_PER_ELEMENT)
+        }
+        await this.gpu.WriteBuffer(this.targetBuffer, new Float32Array(target))
+        const bind = await this.device.createBindGroup({
+            layout: await this.backwardPipline.getBindGroupLayout(0),
+            entries: [
+                {
+                    binding: 0,
+                    resource: {
+                        buffer: probs.mainBuffer
+                    }
+                }, {
+                    binding: 2,
+                    resource: {
+                        buffer: this.targetBuffer
+                    }
+                }, {
+                    binding: 3,
+                    resource: {
+                        buffer: this.paramsForward
+                    }
+                }
+            ]
+        })
+        await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((tokenCount * this.vocabSize) / 256))
+        if (!encoderMain) this.device.queue.submit([encoder.finish()])
+        return probs
+    }
+    async crossEntropy(probs, target, encoderMain) {
+        const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder();
+        const tokenCount = probs.size / (this.vocabSize * Float32Array.BYTES_PER_ELEMENT);
+        if (!Number.isInteger(tokenCount) || tokenCount <= 0) {
+            throw new Error("Invalid probability buffer size or vocabulary size");
+        }
+        const targetBytes = target.length * Uint32Array.BYTES_PER_ELEMENT;
+
+        if (!this.targetBuffer || this.targetBuffer.size < targetBytes) {
+            this.targetBuffer =
+                await this.gpu.createBuffer(targetBytes);
+        }
+
+        if (!this.loss || this.loss.size < tokenCount * 4) {
+            this.loss = await this.gpu.createBuffer(tokenCount * 4);
+        }
+
+        if (!this.acc) {
+            this.acc = await this.gpu.createBuffer(4);
+        }
+
+        if (!this.totalAcc) {
+            this.totalAcc = await this.gpu.createBuffer(4);
+            await this.gpu.WriteBuffer(
+                this.totalAcc,
+                new Float32Array([0])
+            );
+        }
+
+        if (!this.totalLoss) {
+            this.totalLoss = await this.gpu.createBuffer(4);
+            await this.gpu.WriteBuffer(
+                this.totalLoss,
+                new Float32Array([0])
+            );
+        }
+
+        await this.gpu.WriteBuffer(
+            this.targetBuffer,
+            new Uint32Array(target)
+        );
+
+        const bindGroup1 = await this.gpu.createBindGroup(
+            this.lossPipline,
+            0,
+            [probs.mainBuffer, this.targetBuffer, this.loss, this.paramsForward]
+        );
+
+        const bindGroup2 = await this.gpu.createBindGroup(
+            this.accPipline,
+            0,
+            [this.paramsForward, probs.mainBuffer, this.acc, this.targetBuffer]
+        );
+
+        const bindGroup3 = await this.gpu.createBindGroup(
+            this.totalCalculate,
+            0,
+            [this.loss, this.acc, this.totalLoss, this.totalAcc]
+        );
+
+        await this.gpu.RunPipline(
+            encoder,
+            this.lossPipline,
+            bindGroup1,
+            Math.ceil(tokenCount / 4)
+        );
+
+        await this.gpu.RunPipline(
+            encoder,
+            this.accPipline,
+            bindGroup2,
+            1
+        );
+
+        await this.gpu.RunPipline(
+            encoder,
+            this.totalCalculate,
+            bindGroup3,
+            1
+        );
+
+        if (!encoderMain) {
+            await this.gpu.submitQueue(encoder);
+        }
+
+        return {
+            loss: async () =>
+                (await this.gpu.readBuffer(4, this.totalLoss))[0],
+
+            acc: async () =>
+                (await this.gpu.readBuffer(4, this.totalAcc))[0],
+
+            clear: async () => {
+                await this.gpu.WriteBuffer(
+                    this.totalLoss,
+                    new Float32Array([0])
+                );
+
+                await this.gpu.WriteBuffer(
+                    this.totalAcc,
+                    new Float32Array([0])
+                );
+            }
+        };
+    }
+
+}
+

package/src/Tokenizer/index.d.ts

@@ -0,0 +1,12 @@
+export class Tokenizer{
+    vocabs : string[]
+    vocabSize: number
+    IdtoWord : Record<string, string>
+    WordtoID : Record<string , number>
+    encoder(text: string) : {
+        tokens: string[]
+        tokenIDs: number[]
+    }
+    decoder(tokens: number[] | number[][]) : string
+    trainTokenizer(text: string, mergeSize: number, filename: string) : void
+}

package/src/Tokenizer/index.js

@@ -0,0 +1,164 @@
+import fs, { existsSync, readFileSync } from 'fs'
+import { join, dirname, format } from 'path';
+import { fileURLToPath } from 'url';
+
+const __dirname = dirname(fileURLToPath(import.meta.url))
+export class Tokenizer {
+    vocabs = [];
+    vocabSize = 0;
+    IdtoWord = {};
+    WordtoID = {};
+    EOSToken = 2
+    constructor(filename = "") {
+        if (fs.existsSync(filename)) {
+            console.log("Warning: Tokenizer is loaded with a Custom Vocab File!")
+            console.log("")
+            this.vocabs = JSON.parse(readFileSync(filename, "utf-8"))
+        } else {
+            this.vocabs = JSON.parse(readFileSync(join(__dirname, "vocabs.json"), "utf-8"))
+        }
+        this.vocabSize = this.vocabs.length
+        this.IdtoWord = Object.fromEntries(this.vocabs.map((e, idx) => [idx, e]))
+        this.WordtoID = Object.fromEntries(this.vocabs.map((e, idx) => [e, idx]))
+    }
+    encoder(text = "") {
+        const AddSpaces = text.replaceAll(/<EOS>/gi, " <EOS> ").toLocaleLowerCase().replaceAll(/<eos>/g, " <EOS> ").replaceAll(/([,.;'{}!`~|":?])/g, " $1 ");
+        const splitedWords = AddSpaces.split(/\s+/).filter(e => e.length > 0)
+        const output = []
+        for (let a = 0; a < splitedWords.length; a++) {
+            const word = splitedWords[a]
+            if (typeof this.WordtoID[word] !== "undefined") {
+                output.push(word)
+            } else {
+                let max = word.length
+                output.push("$12")
+                while (true) {
+                    const regPatten = new RegExp(`.{0,${max}}`, "g")
+                    const wordsliced = word.match(regPatten).filter((e) => e.length > 0)
+                    const matched = wordsliced.flatMap(e => {
+                        if (this.WordtoID[e] !== undefined) {
+                            return true
+                        }
+                        return false
+
+                    }).reduce((prev, curr) => prev && curr, true)
+                    if (wordsliced.length == 0) output.push("<unk>")
+                    if (matched) {
+                        wordsliced.map((e) => output.push(e))
+                        output.push("$130")
+                        break
+                    } else {
+                        max -= 1
+                        // break
+                    }
+
+                }
+            }
+        }
+        return { tokens: output, tokenIDs: output.map((e) => this.WordtoID[e]) }
+    }
+    decoder(tokens = [1, 2, 3]) {
+        const output = []
+        let isSpace = false
+        // let index = 0
+        for (let token of tokens) {
+            const formated = this.IdtoWord[String(token)]
+            if (formated == "$12") {
+                output[output.length] = ""
+                isSpace = true
+            } else if (formated == "$130") {
+                isSpace = false
+            } else {
+                if (isSpace) {
+                    output[output.length - 1] += formated
+                } else {
+                    output.push(formated)
+                }
+            }
+        }
+        return output.filter(e => e.trim().length > 0).join(" ").replaceAll(/\s+([,.;'{}`!~\"?])/g, "$1")
+    }
+    trainTokenizer(text = "ejrnienieijr", mergeLength = 5, vocabFileName) {
+        console.log("Training The Tokenizer....")
+        const arrText = text.toLowerCase().replace(/\r\n|\n|\r/g, " <EOS> ").toLocaleLowerCase().replaceAll(/<eos>/g, " <EOS> ").replace(/([,.;'{}!`~|":><?])/g, " $1 ").split(/\s+/).map(e => e.trim()).filter(e => e.length > 0)
+        const reg = new RegExp(`.{0,${mergeLength}}`, "g")
+        const prevVocab = existsSync(vocabFileName) ? JSON.parse(readFileSync(vocabFileName, "utf-8")) : [];
+        const specal = ["<unk>",
+            "<PAD>",
+            "<EOS>",
+            "$12",
+            "$130",
+            "😀",
+            "😃",
+            "😄",
+            "😁",
+            "😆",
+            "😅",
+            "😂",
+            "🤣",
+            "😊",
+            "😎",
+            "❤️",
+            "👍",
+            "👎",
+            "🔥",
+            "💯",
+            "🎉",
+            "🚀",
+            "⭐",
+            "😭",
+            "😍",
+            "🥰",
+            "😡",
+            "🤔",
+            "😢",
+            "😮",
+            "👋",
+            "🙏",
+            "💀",
+            "🤖",
+            "👀",
+            "💕",
+            "💔",
+            "😂😂",
+            "😂😭",
+            "🔥🔥",
+            "❤️❤️",
+            "👍👍",
+            "<",
+            ">",
+            ".",
+            "{",
+            "}",
+            "?",
+            "!",
+            "@",
+            "$",
+            ":",
+            "/",
+            "+",
+            "-",
+            "=",
+            "%",
+            "*",
+            "@"
+        ]
+        const uniArr = [...new Set([...arrText, ...specal, ...prevVocab])];
+        const output = []
+        for (let a = 0; a < uniArr.length; a++) {
+            let text = uniArr[a].match(reg).filter((e) => e.length > 0);
+            output.push(...text)
+        }
+        let finalOutput = [...new Set([...output])]
+        fs.writeFileSync(vocabFileName, JSON.stringify(finalOutput, null, 2), "utf-8");
+        console.log("_________________________________")
+        console.log(`New VocabSize : ${finalOutput.length - prevVocab.length}`);
+        console.log("_________________________________")
+        console.log("Operation Completed!")
+        this.vocabs = finalOutput;
+        this.vocabSize = finalOutput.length;
+        this.IdtoWord = Object.fromEntries(this.vocabs.map((e, idx) => [idx, e]))
+        this.WordtoID = Object.fromEntries(this.vocabs.map((e, idx) => [e, idx]))
+
+    }
+}