Ollama: extract contextWindow from num_ctx. Closes #309

Note that from testing, only yarn-mistral has a number set that's not 4096, while some models don't have parameters, don't have a 'num_ctx' value to parse within, or have it set to 4096.
enricoros · Jan 26, 2024 · 18a1629 · 18a1629
1 parent db1346f
commit 18a1629
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/src/modules/llms/server/ollama/ollama.router.ts b/src/modules/llms/server/ollama/ollama.router.ts
@@ -210,6 +210,25 @@ export const llmOllamaRouter = createTRPCRouter({
           const label = capitalizeFirstLetter(modelName) + ((modelTag && modelTag !== 'latest') ? ` · ${modelTag}` : '');
           const description = OLLAMA_BASE_MODELS[modelName]?.description ?? 'Model unknown';
 
+          /* Find the context window from the 'num_ctx' line in the parameters string, if present
+           *  - https://github.com/enricoros/big-AGI/issues/309
+           *  - Note: as of 2024-01-26 the num_ctx line is present in 50% of the models, and in most cases set to 4096
+           *  - We are tracking the Upstream issue https://github.com/ollama/ollama/issues/1473 for better ways to do this in the future
+           */
+          let contextWindow = 4096;
+          if (model.parameters) {
+            // split the parameters into lines, and find one called "num_ctx ...spaces... number"
+            const paramsNumCtx = model.parameters.split('\n').find(line => line.startsWith('num_ctx '));
+            if (paramsNumCtx) {
+              const numCtxValue: string = paramsNumCtx.split(/\s+/)[1];
+              if (numCtxValue) {
+                const numCtxNumber: number = parseInt(numCtxValue);
+                if (!isNaN(numCtxNumber))
+                  contextWindow = numCtxNumber;
+              }
+            }
+          }
+
           // console.log('>>> ollama model', model.name, model.template, model.modelfile, '\n');
 
           return {
@@ -218,7 +237,7 @@ export const llmOllamaRouter = createTRPCRouter({
             created: Date.parse(model.modified_at) ?? undefined,
             updated: Date.parse(model.modified_at) ?? undefined,
             description: description, // description: (model.license ? `License: ${model.license}. Info: ` : '') + model.modelfile || 'Model unknown',
-            contextWindow: 4096, // FIXME: request this information upstream?
+            contextWindow,
             interfaces: [LLM_IF_OAI_Chat],
           } satisfies ModelDescriptionSchema;
         }),

diff --git a/src/modules/llms/server/ollama/ollama.wiretypes.ts b/src/modules/llms/server/ollama/ollama.wiretypes.ts
@@ -1,6 +1,15 @@
 import { z } from 'zod';
 
 
+/*const wireOllamaModelDetailsSchema = z.object({
+  parent_model: z.string().optional(),
+  format: z.string().optional(),                      // e.g. gguf
+  family: z.string().optional(),                      // e.g. llama, phi2, stablelm
+  families: z.array(z.string()).nullable().optional(),// e.g. null, [llama], [phi2]
+  parameter_size: z.string().optional(),              // e.g. 7B, 3B, 34B, 6B, 30B
+  quantization_level: z.string().optional(),          // e.g. Q4_0, Q8_0, F16, ...
+}).optional();*/
+
 /**
  * List Local Models (/api/tags) - Response
  */
@@ -9,7 +18,8 @@ export const wireOllamaListModelsSchema = z.object({
     name: z.string(),
     modified_at: z.string(),
     size: z.number(),
-    digest: z.string(),
+    // digest: z.string(),
+    // details: wireOllamaModelDetailsSchema.optional(),
   })),
 });
 
@@ -21,6 +31,7 @@ export const wireOllamaModelInfoSchema = z.object({
   modelfile: z.string(),
   parameters: z.string().optional(),
   template: z.string().optional(),
+  // details: wireOllamaModelDetailsSchema.optional(),
 });