Selaa lähdekoodia

Add additional options to the generate method.

saul 1 vuosi sitten
vanhempi
commit
5d090f30f5
3 muutettua tiedostoa jossa 53 lisäystä ja 4 poistoa
  1. 6 1
      README.md
  2. 15 3
      src/index.ts
  3. 32 0
      src/interfaces.ts

+ 6 - 1
README.md

@@ -53,11 +53,16 @@ Create a new API handler for ollama.
 ### generate
 
 ```javascript
-ollama.generate(model, prompt);
+ollama.generate(model, prompt, [options]);
 ```
 
 - `model` `<string>` The name of the model to use for the prompt.
 - `prompt` `<string>` The prompt to give the model.
+- `options` `<GenerateOptions>` Optional options to pass to the model.
+  - `parameters` `<ModelParameters>` Model Parameters.
+  - `context` `<number[]>` Context returned from previous calls.
+  - `template` `<string>` Override the default template.
+  - `system` `<string>` Override the default system string.
 - Returns: `<AsyncGenerator<string, GenerateResult>>` A generator that outputs the tokens as strings.
 
 Generate a response for a given prompt with a provided model. The final response object will include statistics and additional data from the request.

+ 15 - 3
src/index.ts

@@ -11,7 +11,9 @@ import type {
 	CreateStatus,
 	PullResponse,
 	PullResult,
-	EmbeddingsResponse
+	EmbeddingsResponse,
+	GenerateOptions,
+	GenerateRequest
 } from "./interfaces.js";
 
 export class Ollama {
@@ -34,8 +36,18 @@ export class Ollama {
 		}));
 	}
 
-	async * generate (model: string, prompt: string): AsyncGenerator<string, GenerateResult> {
-		const response = await utils.post(`${this.config.address}/api/generate`, { model, prompt });
+	async * generate (model: string, prompt: string, options?: Partial<GenerateOptions>): AsyncGenerator<string, GenerateResult> {
+		const parameters = options?.parameters;
+
+		delete options?.parameters;
+
+		const request: GenerateRequest = { model, prompt, ...options };
+
+		if (parameters != null) {
+			request.options = parameters;
+		}
+
+		const response = await utils.post(`${this.config.address}/api/generate`, { ...request });
 
 		if (!response.body) {
 			throw new Error("Missing body");

+ 32 - 0
src/interfaces.ts

@@ -2,6 +2,29 @@ export interface Config {
 	address: string
 }
 
+export interface ModelParameters {
+	mirostat: number
+	mirostat_eta: number
+	mirostat_tau: number
+	num_ctx: number
+	num_gqa: number
+	num_thread: number
+	repeat_last_n: number
+	repeat_penalty: number
+	temperature: number
+	stop: string
+	tfs_z: number
+	top_k: number
+	top_p: number
+}
+
+export interface GenerateOptions {
+	parameters: Partial<ModelParameters>
+	context: number[]
+	template: string
+	system: string
+}
+
 export interface GenerateResult {
 	model: string
 	createdAt: Date
@@ -39,6 +62,15 @@ export interface TagsResponse {
 	}[]
 }
 
+export interface GenerateRequest {
+	model: string
+	prompt: string
+	options?: Partial<ModelParameters>
+	system?: string
+	template?: string
+	context?: number[]
+}
+
 export interface GenerateResponse {
 	model: string
 	created_at: string