inference-server 1.0.0-beta.19

package/src/engines/gpt4all/engine.ts

@@ -0,0 +1,441 @@
+import path from 'node:path'
+import fs from 'node:fs'
+import {
+	loadModel,
+	createCompletion,
+	createEmbedding,
+	InferenceModel,
+	LoadModelOptions,
+	CompletionInput,
+	EmbeddingModel,
+	DEFAULT_MODEL_LIST_URL,
+} from 'gpt4all'
+import {
+	EngineTextCompletionArgs,
+	EngineChatCompletionArgs,
+	EngineChatCompletionResult,
+	EngineTextCompletionResult,
+	CompletionFinishReason,
+	EngineContext,
+	EngineEmbeddingArgs,
+	EngineEmbeddingResult,
+	FileDownloadProgress,
+	ModelConfig,
+	TextCompletionParams,
+	ChatMessage,
+} from '#package/types/index.js'
+import { LogLevels } from '#package/lib/logger.js'
+import { downloadModelFile } from '#package/lib/downloadModelFile.js'
+import { acquireFileLock } from '#package/lib/acquireFileLock.js'
+import { validateModelFile } from '#package/lib/validateModelFile.js'
+import { createChatMessageArray } from './util.js'
+
+export type GPT4AllInstance = InferenceModel | EmbeddingModel
+
+export interface GPT4AllModelMeta {
+	url: string
+	md5sum: string
+	filename: string
+	promptTemplate: string
+	systemPrompt: string
+	filesize: number
+	ramrequired: number
+}
+
+export interface GPT4AllModelConfig extends ModelConfig {
+	location: string
+	md5?: string
+	url?: string
+	contextSize?: number
+	batchSize?: number
+	task: 'text-completion' | 'embedding'
+	initialMessages?: ChatMessage[]
+	completionDefaults?: TextCompletionParams
+	device?: {
+		gpu?: boolean | 'auto' | (string & {})
+		gpuLayers?: number
+		cpuThreads?: number
+	}
+}
+
+export const autoGpu = true
+
+export async function prepareModel(
+	{ config, log }: EngineContext<GPT4AllModelConfig>,
+	onProgress?: (progress: FileDownloadProgress) => void,
+	signal?: AbortSignal,
+) {
+	fs.mkdirSync(path.dirname(config.location), { recursive: true })
+	const releaseFileLock = await acquireFileLock(config.location)
+	if (signal?.aborted) {
+		releaseFileLock()
+		return
+	}
+	log(LogLevels.info, `Preparing gpt4all model at ${config.location}`, {
+		model: config.id,
+	})
+	let modelMeta: GPT4AllModelMeta | undefined
+	let modelList: GPT4AllModelMeta[]
+	const modelMetaPath = path.join(path.dirname(config.location), 'models.json')
+	try {
+		if (!fs.existsSync(modelMetaPath)) {
+			const res = await fetch(DEFAULT_MODEL_LIST_URL)
+			modelList = (await res.json()) as GPT4AllModelMeta[]
+			fs.writeFileSync(modelMetaPath, JSON.stringify(modelList, null, 2))
+		} else {
+			modelList = JSON.parse(fs.readFileSync(modelMetaPath, 'utf-8'))
+		}
+		const foundModelMeta = modelList.find((item) => {
+			if (config.md5 && item.md5sum) {
+				return item.md5sum === config.md5
+			}
+			if (config.url && item.url) {
+				return item.url === config.url
+			}
+			return item.filename === path.basename(config.location)
+		})
+		if (foundModelMeta) {
+			modelMeta = foundModelMeta
+		}
+
+		const validationError = await validateModelFile({
+			...config,
+			md5: config.md5 || modelMeta?.md5sum,
+		})
+		if (signal?.aborted) {
+			return
+		}
+		if (validationError) {
+			if (config.url) {
+				log(LogLevels.info, 'Downloading', {
+					model: config.id,
+					url: config.url,
+					location: config.location,
+					error: validationError,
+				})
+				await downloadModelFile({
+					url: config.url,
+					filePath: config.location,
+					modelsCachePath: config.modelsCachePath,
+					onProgress,
+					signal,
+				})
+			} else {
+				throw new Error(`${validationError} - No URL provided`)
+			}
+		}
+
+		const finalValidationError = await validateModelFile({
+			...config,
+			md5: config.md5 || modelMeta?.md5sum,
+		})
+		if (finalValidationError) {
+			throw new Error(`Downloaded files are invalid: ${finalValidationError}`)
+		}
+		if (signal?.aborted) {
+			return
+		}
+
+		return modelMeta
+	} catch (error) {
+		throw error
+	} finally {
+		releaseFileLock()
+	}
+}
+
+export async function createInstance({ config, log }: EngineContext<GPT4AllModelConfig>, signal?: AbortSignal) {
+	log(LogLevels.info, `Load GPT4All model ${config.location}`)
+	let device = config.device?.gpu ?? 'cpu'
+	if (typeof device === 'boolean') {
+		device = device ? 'gpu' : 'cpu'
+	} else if (device === 'auto') {
+		device = 'cpu'
+	}
+	const loadOpts: LoadModelOptions = {
+		modelPath: path.dirname(config.location),
+		// file: config.file,
+		modelConfigFile: path.dirname(config.location) + '/models.json',
+		allowDownload: false,
+		device: device,
+		ngl: config.device?.gpuLayers ?? 100,
+		nCtx: config.contextSize ?? 2048,
+		// verbose: true,
+		// signal?: // TODO no way to cancel load
+	}
+
+	let modelType: 'inference' | 'embedding'
+	if (config.task === 'text-completion') {
+		modelType = 'inference'
+	} else if (config.task === 'embedding') {
+		modelType = 'embedding'
+	} else {
+		throw new Error(`Unsupported task type: ${config.task}`)
+	}
+
+	const instance = await loadModel(path.basename(config.location), {
+		...loadOpts,
+		type: modelType,
+	})
+	if (config.device?.cpuThreads) {
+		instance.llm.setThreadCount(config.device.cpuThreads)
+	}
+
+	if ('generate' in instance) {
+		if (config.initialMessages?.length) {
+			let messages = createChatMessageArray(config.initialMessages)
+			let systemPrompt
+			if (messages[0].role === 'system') {
+				systemPrompt = messages[0].content
+				messages = messages.slice(1)
+			}
+			await instance.createChatSession({
+				systemPrompt,
+				messages,
+			})
+		} else if (config.prefix) {
+			await instance.generate(config.prefix, {
+				nPredict: 0,
+			})
+		} else {
+			await instance.generate('', {
+				nPredict: 0,
+			})
+		}
+	}
+
+	return instance
+}
+
+export async function disposeInstance(instance: GPT4AllInstance) {
+	instance.dispose()
+}
+
+export async function processTextCompletionTask(
+	{ request, config, onChunk }: EngineTextCompletionArgs<GPT4AllModelConfig>,
+	instance: GPT4AllInstance,
+	signal?: AbortSignal,
+): Promise<EngineTextCompletionResult> {
+	if (!('generate' in instance)) {
+		throw new Error('Instance does not support text completion.')
+	}
+	if (!request.prompt) {
+		throw new Error('Prompt is required for text completion.')
+	}
+
+	let finishReason: CompletionFinishReason = 'eogToken'
+	let suffixToRemove: string | undefined
+
+	const defaults = config.completionDefaults ?? {}
+	const stopTriggers = request.stop ?? defaults.stop ?? []
+	const includesStopTriggers = (text: string) => stopTriggers.find((t) => text.includes(t))
+	const result = await instance.generate(request.prompt, {
+		// @ts-ignore
+		special: true, // allows passing in raw prompt (including <|start|> etc.)
+		promptTemplate: '%1',
+		temperature: request.temperature ?? defaults.temperature,
+		nPredict: request.maxTokens ?? defaults.maxTokens,
+		topP: request.topP ?? defaults.topP,
+		topK: request.topK ?? defaults.topK,
+		minP: request.minP ?? defaults.minP,
+		nBatch: config?.batchSize,
+		repeatLastN: request.repeatPenaltyNum ?? defaults.repeatPenaltyNum,
+		// repeat penalty is doing something different than both frequency and presence penalty
+		// so not falling back to them here.
+		repeatPenalty: request.repeatPenalty ?? defaults.repeatPenalty,
+		// seed: args.seed, // https://github.com/nomic-ai/gpt4all/issues/1952
+		// @ts-ignore
+		onResponseToken: (tokenId, text) => {
+			const matchingTrigger = includesStopTriggers(text)
+			if (matchingTrigger) {
+				finishReason = 'stopTrigger'
+				suffixToRemove = text
+				return false
+			}
+			if (onChunk) {
+				onChunk({
+					text,
+					tokens: [tokenId],
+				})
+			}
+			return !signal?.aborted
+		},
+		// @ts-ignore
+		onResponseTokens: ({ tokenIds, text }) => {
+			const matchingTrigger = includesStopTriggers(text)
+			if (matchingTrigger) {
+				finishReason = 'stopTrigger'
+				suffixToRemove = text
+				return false
+			}
+			if (onChunk) {
+				onChunk({
+					text,
+					tokens: tokenIds,
+				})
+			}
+			return !signal?.aborted
+		},
+	})
+
+	if (result.tokensGenerated === request.maxTokens) {
+		finishReason = 'maxTokens'
+	}
+
+	let responseText = result.text
+	if (suffixToRemove) {
+		responseText = responseText.slice(0, -suffixToRemove.length)
+	}
+
+	return {
+		finishReason,
+		text: responseText,
+		promptTokens: result.tokensIngested,
+		completionTokens: result.tokensGenerated,
+		contextTokens: instance.activeChatSession?.promptContext.nPast ?? 0,
+	}
+}
+
+export async function processChatCompletionTask(
+	{ request, config, resetContext, log, onChunk }: EngineChatCompletionArgs<GPT4AllModelConfig>,
+	instance: GPT4AllInstance,
+	signal?: AbortSignal,
+): Promise<EngineChatCompletionResult> {
+	if (!('createChatSession' in instance)) {
+		throw new Error('Instance does not support chat completion.')
+	}
+	let session = instance.activeChatSession
+	if (!session || resetContext) {
+		log(LogLevels.debug, 'Resetting chat context')
+		let messages = createChatMessageArray(request.messages)
+		let systemPrompt
+		if (messages[0].role === 'system') {
+			systemPrompt = messages[0].content
+			messages = messages.slice(1)
+		}
+		// drop last user message
+		if (messages[messages.length - 1].role === 'user') {
+			messages = messages.slice(0, -1)
+		}
+
+		session = await instance.createChatSession({
+			systemPrompt,
+			messages,
+		})
+	}
+
+	const conversationMessages = createChatMessageArray(request.messages).filter((m) => m.role !== 'system')
+
+	const lastMessage = conversationMessages[conversationMessages.length - 1]
+	if (!(lastMessage.role === 'user' && lastMessage.content)) {
+		throw new Error('Chat completions require a final user message.')
+	}
+	const input: CompletionInput = lastMessage.content
+
+	let finishReason: CompletionFinishReason = 'eogToken'
+	let suffixToRemove: string | undefined
+
+	const defaults = config.completionDefaults ?? {}
+	const stopTriggers = request.stop ?? defaults.stop ?? []
+	const includesStopTriggers = (text: string) => stopTriggers.find((t) => text.includes(t))
+	const result = await createCompletion(session, input, {
+		temperature: request.temperature ?? defaults.temperature,
+		nPredict: request.maxTokens ?? defaults.maxTokens,
+		topP: request.topP ?? defaults.topP,
+		topK: request.topK ?? defaults.topK,
+		minP: request.minP ?? defaults.minP,
+		nBatch: config.batchSize,
+		repeatLastN: request.repeatPenaltyNum ?? defaults.repeatPenaltyNum,
+		repeatPenalty: request.repeatPenalty ?? defaults.repeatPenalty,
+		// seed: args.seed, // see https://github.com/nomic-ai/gpt4all/issues/1952
+		// @ts-ignore
+		onResponseToken: (tokenId, text) => {
+			const matchingTrigger = includesStopTriggers(text)
+			if (matchingTrigger) {
+				finishReason = 'stopTrigger'
+				suffixToRemove = text
+				return false
+			}
+			if (onChunk) {
+				onChunk({
+					text,
+					tokens: [tokenId],
+				})
+			}
+			return !signal?.aborted
+		},
+		// @ts-ignore
+		onResponseTokens: ({ tokenIds, text }) => {
+			const matchingTrigger = includesStopTriggers(text)
+			if (matchingTrigger) {
+				finishReason = 'stopTrigger'
+				suffixToRemove = text
+				return false
+			}
+			if (onChunk) {
+				onChunk({
+					tokens: tokenIds,
+					text,
+				})
+			}
+
+			return !signal?.aborted
+		},
+	})
+
+	if (result.usage.completion_tokens === request.maxTokens) {
+		finishReason = 'maxTokens'
+	}
+
+	let response = result.choices[0].message.content
+	if (suffixToRemove) {
+		response = response.slice(0, -suffixToRemove.length)
+	}
+
+	return {
+		finishReason,
+		message: {
+			role: 'assistant',
+			content: response,
+		},
+		promptTokens: result.usage.prompt_tokens,
+		completionTokens: result.usage.completion_tokens,
+		contextTokens: session.promptContext.nPast,
+	}
+}
+
+export async function processEmbeddingTask(
+	{ request, config }: EngineEmbeddingArgs,
+	instance: GPT4AllInstance,
+	signal?: AbortSignal,
+): Promise<EngineEmbeddingResult> {
+	if (!('embed' in instance)) {
+		throw new Error('Instance does not support embedding.')
+	}
+	if (!request.input) {
+		throw new Error('Input is required for embedding.')
+	}
+	const texts: string[] = []
+	if (typeof request.input === 'string') {
+		texts.push(request.input)
+	} else if (Array.isArray(request.input)) {
+		for (const input of request.input) {
+			if (typeof input === 'string') {
+				texts.push(input)
+			} else if (input.type === 'text') {
+				texts.push(input.content)
+			} else if (input.type === 'image') {
+				throw new Error('Image inputs not implemented.')
+			}
+		}
+	}
+
+	const res = await createEmbedding(instance, texts, {
+		dimensionality: request.dimensions,
+	})
+
+	return {
+		embeddings: res.embeddings,
+		inputTokens: res.n_prompt_tokens,
+	}
+}

package/src/engines/gpt4all/util.ts

@@ -0,0 +1,31 @@
+import { ChatMessage as GPT4AllChatMessage } from 'gpt4all'
+import { ChatMessage } from '#package/types/index.js'
+import { flattenMessageTextContent } from '#package/lib/flattenMessageTextContent.js'
+
+export function createChatMessageArray(
+	messages: ChatMessage[],
+): GPT4AllChatMessage[] {
+	const chatMessages: GPT4AllChatMessage[] = []
+	let systemPrompt: string | undefined
+	for (const message of messages) {
+		if (message.role === 'user' || message.role === 'assistant') {
+			chatMessages.push({
+				role: message.role,
+				content: flattenMessageTextContent(message.content),
+			})
+		} else if (message.role === 'system') {
+			if (systemPrompt) {
+				systemPrompt += '\n\n' + message.content
+			} else {
+				systemPrompt = flattenMessageTextContent(message.content)
+			}
+		}
+	}
+	if (systemPrompt) {
+		chatMessages.unshift({
+			role: 'system',
+			content: systemPrompt,
+		})
+	}
+	return chatMessages
+}

package/src/engines/index.ts

@@ -0,0 +1,28 @@
+import type { ModelPool } from '#package/pool.js'
+import type { ModelStore } from '#package/store.js'
+import { ModelEngine, EngineStartContext, ModelOptions, BuiltInModelOptions } from '#package/types/index.js'
+
+export const BuiltInEngines = {
+	gpt4all: 'gpt4all',
+	nodeLlamaCpp: 'node-llama-cpp',
+	transformersJs: 'transformers-js',
+	stableDiffusionCpp: 'stable-diffusion-cpp',
+} as const
+
+export type BuiltInEngineName = typeof BuiltInEngines[keyof typeof BuiltInEngines];
+
+export const builtInEngineNames: string[] = [
+	...Object.values(BuiltInEngines),
+] as const
+
+export class CustomEngine implements ModelEngine {
+	pool!: ModelPool
+	store!: ModelStore
+	async start({ pool, store }: EngineStartContext) {
+		this.pool = pool
+		this.store = store
+	}
+	async prepareModel() {}
+	async createInstance() {}
+	async disposeInstance() {}
+}