Upload serve_mamba.js with huggingface_hub

serve_mamba.js

@@ -0,0 +1,86 @@
+const http = require('http'), fs = require('fs'), path = require('path');
+const PORT = 8140;
+const ROOT = __dirname;
+const HF = path.join(process.env.USERPROFILE, '.cache', 'huggingface', 'hub',
+  'models--tiiuae--falcon-mamba-7b-instruct', 'snapshots',
+  'b250fc9399d14f56aca18e9ea70bbfb1f73479eb');
+
+const MIME = {
+  '.html': 'text/html', '.js': 'text/javascript', '.wgsl': 'text/plain',
+  '.json': 'application/json', '.safetensors': 'application/octet-stream'
+};
+
+const { execFile } = require('child_process');
+
+function handleAPI(req, res) {
+  let body = '';
+  req.on('data', c => body += c);
+  req.on('end', () => {
+    const data = JSON.parse(body);
+    const p = req.url;
+    const script = p === '/tokenize'
+      ? `import sys,json;from transformers import AutoTokenizer;t=AutoTokenizer.from_pretrained('tiiuae/falcon-mamba-7b-instruct');d=json.loads(sys.stdin.read());text=d['text'];msgs=[{'role':'user','content':text}];templated=t.apply_chat_template(msgs,tokenize=False,add_generation_prompt=True);print(json.dumps(t.encode(templated)))`
+      : `import sys,json;from transformers import AutoTokenizer;t=AutoTokenizer.from_pretrained('tiiuae/falcon-mamba-7b-instruct');print(json.dumps(t.decode(json.loads(sys.stdin.read())['tokens'])))`;
+    const child = execFile('python', ['-c', script], { maxBuffer: 10*1024*1024 }, (err, stdout) => {
+      res.writeHead(200, { 'Content-Type': 'application/json', 'Access-Control-Allow-Origin': '*' });
+      if (err) { res.end(JSON.stringify({ error: err.message })); return; }
+      res.end(JSON.stringify({ result: JSON.parse(stdout.trim()) }));
+    });
+    child.stdin.write(JSON.stringify(data));
+    child.stdin.end();
+  });
+}
+
+http.createServer((req, res) => {
+  if (req.method === 'OPTIONS') {
+    res.writeHead(200, { 'Access-Control-Allow-Origin': '*', 'Access-Control-Allow-Methods': 'GET,POST', 'Access-Control-Allow-Headers': 'Content-Type' });
+    return res.end();
+  }
+  if (req.method === 'POST' && (req.url === '/tokenize' || req.url === '/detokenize')) {
+    return handleAPI(req, res);
+  }
+
+  let p = decodeURIComponent(req.url.split('?')[0]);
+  if (p === '/') p = '/index.html';
+
+  let fp;
+  if (p.startsWith('/weights/')) {
+    fp = path.join(HF, p.slice('/weights/'.length));
+  } else {
+    fp = path.join(ROOT, p);
+  }
+
+  fs.stat(fp, (e, st) => {
+    if (e) { res.writeHead(404); return res.end('not found: ' + fp); }
+    const total = st.size;
+    const range = req.headers.range;
+    const ct = MIME[path.extname(fp)] || 'application/octet-stream';
+
+    if (range) {
+      const m = range.match(/bytes=(\d+)-(\d*)/);
+      if (m) {
+        const start = parseInt(m[1]);
+        const end = m[2] ? parseInt(m[2]) : total - 1;
+        res.writeHead(206, {
+          'Content-Type': ct,
+          'Content-Range': 'bytes ' + start + '-' + end + '/' + total,
+          'Content-Length': end - start + 1,
+          'Accept-Ranges': 'bytes',
+          'Access-Control-Allow-Origin': '*'
+        });
+        fs.createReadStream(fp, { start, end }).pipe(res);
+        return;
+      }
+    }
+
+    res.writeHead(200, {
+      'Content-Type': ct, 'Content-Length': total,
+      'Accept-Ranges': 'bytes', 'Access-Control-Allow-Origin': '*'
+    });
+    fs.createReadStream(fp).pipe(res);
+  });
+}).listen(PORT, () => {
+  console.log('Mamba WebGPU on :' + PORT);
+  console.log('Weights from: ' + HF);
+  console.log('HF exists: ' + fs.existsSync(HF));
+});