| #!/usr/bin/env node |
| |
| |
| |
| |
| |
| |
| |
| |
| 'use strict'; |
| const fs = require('fs'); |
| const path = require('path'); |
| const http = require('http'); |
|
|
| const HEAD = JSON.parse(fs.readFileSync(path.join(__dirname, '..', 'ear', 'head.json'), 'utf8')); |
| const BASE = process.env.FLM_BASE || 'http://localhost:52625'; |
|
|
| function post(urlStr, body) { |
| return new Promise((resolve, reject) => { |
| const u = new URL(urlStr); |
| const data = JSON.stringify(body); |
| const req = http.request({ hostname: u.hostname, port: u.port, path: u.pathname, method: 'POST', |
| headers: { 'Content-Type': 'application/json', 'Content-Length': Buffer.byteLength(data) }, timeout: 60000 }, |
| res => { let b = ''; res.on('data', c => b += c); res.on('end', () => { try { resolve(JSON.parse(b)); } catch (e) { reject(new Error('bad json: ' + b.slice(0, 200))); } }); }); |
| req.on('error', reject); |
| req.on('timeout', () => { req.destroy(); reject(new Error('timeout')); }); |
| req.write(data); req.end(); |
| }); |
| } |
|
|
| async function embed(text) { |
| let r; |
| try { |
| r = await post(BASE + '/v1/embeddings', { model: 'embed-gemma:300m', input: text }); |
| } catch (e) { |
| if (/ECONNRESET|socket hang up/i.test(e.message)) { |
| await new Promise(res => setTimeout(res, 400)); |
| r = await post(BASE + '/v1/embeddings', { model: 'embed-gemma:300m', input: text }); |
| } else throw e; |
| } |
| const v = r.data && r.data[0] && r.data[0].embedding; |
| if (!v) throw new Error('no embedding in response: ' + JSON.stringify(r).slice(0, 200)); |
| |
| const n = Math.sqrt(v.reduce((s, x) => s + x * x, 0)) || 1; |
| return v.map(x => x / n); |
| } |
|
|
| function applyHead(v) { |
| const { W, b, labels } = HEAD; |
| const logits = b.slice(); |
| for (let i = 0; i < v.length; i++) { |
| const wi = W[i]; |
| for (let j = 0; j < logits.length; j++) logits[j] += v[i] * wi[j]; |
| } |
| const m = Math.max(...logits); |
| const exps = logits.map(x => Math.exp(x - m)); |
| const Z = exps.reduce((a, x) => a + x, 0); |
| const probs = exps.map(x => x / Z); |
| const top = probs.indexOf(Math.max(...probs)); |
| return { label: labels[top], confidence: +probs[top].toFixed(3), |
| all: Object.fromEntries(labels.map((l, j) => [l, +probs[j].toFixed(3)])) }; |
| } |
|
|
| async function identify(text) { |
| const t0 = Date.now(); |
| const v = await embed(text); |
| const out = applyHead(v); |
| out.ms = Date.now() - t0; |
| return out; |
| } |
|
|
| module.exports = { identify, applyHead, embed }; |
|
|
| if (require.main === module) { |
| (async () => { |
| const arg = process.argv.slice(2).join(' '); |
| if (arg === '--server-check') { |
| try { await embed('hello'); console.log('embed endpoint OK at ' + BASE); } |
| catch (e) { console.log('embed endpoint unavailable: ' + e.message.slice(0, 120)); console.log('start one with: flm serve embed-gemma:300m --embed 1 (or flm serve <chat-model> --embed 1)'); } |
| return; |
| } |
| if (!arg) { console.error('usage: node bin/ear.js "text" | --server-check'); process.exit(1); } |
| try { |
| console.log(JSON.stringify(await identify(arg), null, 2)); |
| } catch (e) { |
| console.error('ear error: ' + e.message); |
| process.exit(1); |
| } |
| })(); |
| } |
|
|