thot-experiment · blueskymonster · Mar 28, 2024 · Apr 18, 2024
diff --git a/LICENSE b/LICENSE
@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2023 thot-experiment
+Copyright (c) 2024 Hallucinate, LLC
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

diff --git a/README.md b/README.md
@@ -1,25 +1,23 @@
-# oobapi-ws
-[oobabooga/text-generation-webui](https://github.com/oobabooga/text-generation-webui) websocket api wrapper
+# oobapi-stream
+[oobabooga/text-generation-webui](https://github.com/oobabooga/text-generation-webui) streaming api wrapper
 
-  This is a simple isomorphic wrapper for the oobabooga websocket api (must be enabled when [launching ooba](https://github.com/oobabooga/text-generation-webui#api) with `--api`)
+  This is a simple isomorphic wrapper for the oobabooga api that handles the necessary HTTP shenanigans to stream token predictions to you as they come instead of getting a result all at once. To use it, API mode must be enabled when [launching ooba](https://github.com/oobabooga/text-generation-webui#api) by using `--api`
 
 ## Usage
 
-  As this wrapper is meant to function the same in the browser as well as in node it must be passed a `WebSocket` object when instantiated. This can be either the browser context `WebSocket` or one from the [`ws` module](https://github.com/websockets/ws). See `node_example.mjs` and `browser_example.html` for more details.
-
   Ｗhen instantiating the api you can specify a host and a port, as well as default generation options．
   ```js
-  import ooba from 'oobapi-ws'
+  import ooba from 'oobapi-stream'
   const generation_options = {max_new_tokens: 1000}
   const host = 'remote.host.net'
   const port = 1337
-  const api = await ooba(WebSocket)({host,port}, generation_options)
+  const api = ooba({host,port}, generation_options)
   ```
   or if you're running ooba locally and you want to use the default options you can omit both
   ```js
-  const api = await ooba(WebSocket)()
+  const api = ooba()
   ```
-  this promise will return once a websocket connection to `localhost:5005` is established
+  this will return an api object that is ready to make streaming requests.
 
   in order to generate text simply call `api.generate()` with a prompt
   ```js

diff --git a/browser_example.html b/browser_example.html
@@ -11,9 +11,9 @@
   </head>
   <body>
     <script type="module">
-      import ooba from './oobapi-ws.mjs'
+      import ooba from './oobapi-stream.mjs'
 
-      const api = await ooba(WebSocket)()
+      const api = ooba()
       const prompt = "In order to make homemade bread, follow these steps:\n1)"
       document.body.innerText += prompt
 

diff --git a/node_example.mjs b/node_example.mjs
@@ -1,11 +1,19 @@
-import WebSocket from 'ws'
-import ooba from './oobapi-ws.mjs'
+import ooba from './oobapi-stream.mjs'
 
-const api = await ooba(WebSocket)()
+const api = ooba()
 
-const prompt = "In order to make homemade bread, follow these steps:\n1)"
+const prompt = "[INST]In order to make homemade bread, follow these steps:[/INST]\n)"
 
+let response_complete = false
 process.stdout.write(prompt)
 api.ontoken = token => process.stdout.write(token)
+api.onend = () => response_complete = true
+api.onerror = () => response_complete = true
 
 api.generate(prompt)
+
+const delay = ms => new Promise(res => setTimeout(res, ms))
+
+while (!response_complete) {
+  await delay(100)
+}
diff --git a/oobapi-stream.mjs b/oobapi-stream.mjs
@@ -0,0 +1,109 @@
+import http from 'http'
+// default options (model, prompt, and stream excepted) from OpenAI API documentation
+const default_options = {
+  model: "unsloth_mistral-7b-instruct-v0.2-bnb-4bit",
+  prompt: "[INST]In five sentences or less, explain what buffer solutions are and how they are relevant to biology.[/INST]",
+  best_of: 1,
+  echo: false,
+  frequency_penalty: 0,
+  logit_bias: {},
+  logprobs: 0,
+  max_tokens: 250,
+  n: 1,
+  presence_penalty: 0,
+  stream: true, // Do not overwrite this option.
+  suffix: "",
+  temperature: 1,
+  top_p: 1,
+  min_p: 0,
+  dynamic_temperature: false,
+  dynatemp_low: 1,
+  dynatemp_high: 1,
+  dynatemp_exponent: 1,
+  smoothing_factor: 0,
+  smoothing_curve: 1,
+  stop: [],
+  top_k: 0,
+  repetition_penalty: 1,
+  repetition_penalty_range: 1024,
+  typical_p: 1,
+  tfs: 1,
+  top_a: 0,
+  epsilon_cutoff: 0,
+  eta_cutoff: 0,
+  guidance_scale: 1,
+  negative_prompt: "",
+  penalty_alpha: 0,
+  mirostat_mode: 0,
+  mirostat_tau: 5,
+  mirostat_eta: 0.1,
+  temperature_last: false,
+  do_sample: true,
+  seed: -1,
+  encoder_repetition_penalty: 1,
+  no_repeat_ngram_size: 0,
+  min_length: 0,
+  num_beams: 1,
+  length_penalty: 1,
+  early_stopping: false,
+  truncation_length: 0,
+  max_tokens_second: 0,
+  prompt_lookup_num_tokens: 0,
+  custom_token_bans: "",
+  sampler_priority: [],
+  ban_eos_token: false,
+  add_bos_token: true,
+  skip_special_tokens: true,
+  grammar_string: ""
+}
+
+const api = ({host='127.0.0.1', port=5000}={}, parameters=default_options) => {
+  const URI = `http://${host}:${port}/v1/completions`
+
+
+  let api = {parameters}
+  api.generate = prompt => {
+    // check for string prompt and wrap in an object
+    if (typeof prompt === 'string') prompt = {prompt}
+    const postData = JSON.stringify(Object.assign({},api.parameters, prompt))
+    const http_options = {
+      hostname: host,
+      port,
+      path: '/v1/completions',
+      method: 'POST',
+      headers: {
+        'Accept': 'application/json',
+        'Content-Type': 'application/json',
+        'Content-Length': Buffer.byteLength(postData)
+      }
+    }
+
+    const req = http.request(http_options, (res) => {
+      res.setEncoding('utf8')
+
+      res.on('data', (chunk) => {
+        // Remove the 'data:' prefix from the chunk
+        const parsed = JSON.parse(chunk.slice(5))
+        // TODO: handle multiples choices when n != 1
+        const response = parsed.choices[0]
+        if (response.finish_reason == null) {
+          api.ontoken?.(response.text)
+          api.debugstream?.(parsed)
+        } else {
+          api.onend?.(parsed)
+        }
+      })
+    })
+
+    req.on('error', (error) => {
+      api.onerror?.({error})
+      console.error("Request error:", error)
+    })
+
+    req.write(postData)
+    req.end()
+  }
+  return api
+}
+
+export default api
diff --git a/oobapi-ws.mjs b/oobapi-ws.mjs
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,14 +1,14 @@
 {
-  "name": "oobapi-ws",
-  "version": "1.0.1",
-  "description": "oobabooga/text-generation-webui websocket api wrapper",
-  "main": "oobapi-ws.mjs",
+  "name": "oobapi-stream",
+  "version": "1.0.3",
+  "description": "oobabooga/text-generation-webui streaming api wrapper",
+  "main": "oobapi-stream.mjs",
   "scripts": {
     "test": "echo \"Error: no test specified\" && exit 1"
   },
   "repository": {
     "type": "git",
-    "url": "git+https://github.com/thot-experiment/oobapi-ws.git"
+    "url": "git+https://github.com/hallucinate-games/oobapi-stream.git"
   },
   "keywords": [
     "oobabooga",
@@ -21,15 +21,24 @@
     "alpaca",
     "large",
     "language",
-    "model"
+    "model",
+    "stream"
+  ],
+  "contributors": [
+    {
+      "name": "Matt Vana",
+      "email": "thot@hallucinate.games",
+      "url": "https://github.com/thot-experiment"
+    },
+    {
+      "name": "Elliot Corvinova",
+      "email": "elliot@hallucinate.games",
+      "url": "https://github.com/blueskymonster"
+    }
   ],
-  "author": "thot-experiment",
   "license": "MIT",
-  "devDependencies": {
-    "ws": "^8.13.0"
-  },
   "bugs": {
-    "url": "https://github.com/thot-experiment/oobapi-ws/issues"
+    "url": "https://github.com/hallucinate-games/oobapi-stream/issues"
   },
-  "homepage": "https://github.com/thot-experiment/oobapi-ws#readme"
+  "homepage": "https://github.com/hallucinate-games/oobapi-streams#readme"
 }