chat.mjs 2.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  1. import * as readline from 'node:readline'
  2. import { stdin, stdout } from 'node:process'
  3. const API_URL = 'http://127.0.0.1:8080'
  4. const chat = [
  5. {
  6. human: "Hello, Assistant.",
  7. assistant: "Hello. How may I help you today?"
  8. },
  9. {
  10. human: "Please tell me the largest city in Europe.",
  11. assistant: "Sure. The largest city in Europe is Moscow, the capital of Russia."
  12. },
  13. ]
  14. const instruction = `A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.`
  15. function format_prompt(question) {
  16. return `${instruction}\n${
  17. chat.map(m =>`### Human: ${m.human}\n### Assistant: ${m.assistant}`).join("\n")
  18. }\n### Human: ${question}\n### Assistant:`
  19. }
  20. async function tokenize(content) {
  21. const result = await fetch(`${API_URL}/tokenize`, {
  22. method: 'POST',
  23. body: JSON.stringify({ content })
  24. })
  25. if (!result.ok) {
  26. return []
  27. }
  28. return await result.json().tokens
  29. }
  30. const n_keep = await tokenize(instruction).length
  31. async function chat_completion(question) {
  32. const result = await fetch(`${API_URL}/completion`, {
  33. method: 'POST',
  34. body: JSON.stringify({
  35. prompt: format_prompt(question),
  36. temperature: 0.2,
  37. top_k: 40,
  38. top_p: 0.9,
  39. n_keep: n_keep,
  40. n_predict: 256,
  41. stop: ["\n### Human:"], // stop completion after generating this
  42. stream: true,
  43. })
  44. })
  45. if (!result.ok) {
  46. return
  47. }
  48. let answer = ''
  49. for await (var chunk of result.body) {
  50. const t = Buffer.from(chunk).toString('utf8')
  51. if (t.startsWith('data: ')) {
  52. const message = JSON.parse(t.substring(6))
  53. answer += message.content
  54. process.stdout.write(message.content)
  55. if (message.stop) {
  56. if (message.truncated) {
  57. chat.shift()
  58. }
  59. break
  60. }
  61. }
  62. }
  63. process.stdout.write('\n')
  64. chat.push({ human: question, assistant: answer.trimStart() })
  65. }
  66. const rl = readline.createInterface({ input: stdin, output: stdout });
  67. const readlineQuestion = (rl, query, options) => new Promise((resolve, reject) => {
  68. rl.question(query, options, resolve)
  69. });
  70. while(true) {
  71. const question = await readlineQuestion(rl, '> ')
  72. await chat_completion(question)
  73. }