Switch chat model back to mistral, use zephyr for small tasks (#515)

* Switch chat model back to mistral, use zephyr for small tasks

* typo

* fix tests

.env.template

@@ -94,24 +94,24 @@ MODELS=`[
       ]
     },
     {
-      "name": "HuggingFaceH4/zephyr-7b-alpha",
-      "displayName": "HuggingFaceH4/zephyr-7b-alpha",
-      "description": "Zephyr 7B α is a fine-tune of Mistral 7B, released by the Hugging Face H4 RLHF team.",
-      "websiteUrl": "https://huggingface.co/HuggingFaceH4/zephyr-7b-alpha/",
+      "name": "mistralai/Mistral-7B-Instruct-v0.1",
+      "displayName": "mistralai/Mistral-7B-Instruct-v0.1",
+      "description": "Mistral 7B is a new Apache 2.0 model, released by Mistral AI that outperforms Llama2 13B in benchmarks.",
+      "websiteUrl": "https://mistral.ai/news/announcing-mistral-7b/",
       "preprompt": "",
-      "chatPromptTemplate" : "<|system|>\n{{preprompt}}</s>\n{{#each messages}}{{#ifUser}}<|user|>\n{{content}}</s>\n<|assistant|>\n{{/ifUser}}{{#ifAssistant}}{{content}}</s>\n{{/ifAssistant}}{{/each}}",
+      "chatPromptTemplate" : "<s>{{#each messages}}{{#ifUser}}[INST] {{#if @first}}{{#if @root.preprompt}}{{@root.preprompt}}\n{{/if}}{{/if}}{{content}} [/INST]{{/ifUser}}{{#ifAssistant}}{{content}}</s>{{/ifAssistant}}{{/each}}",
       "parameters": {
-        "temperature": 0.7,
+        "temperature": 0.1,
         "top_p": 0.95,
         "repetition_penalty": 1.2,
         "top_k": 50,
         "truncate": 1000,
         "max_new_tokens": 2048,
-        "stop": ["</s>", "<|>"]
+        "stop": ["</s>"]
       },
       "promptExamples": [
       {
-         "title": "Write an email from bullet list",
+          "title": "Write an email from bullet list",
         "prompt": "As a restaurant owner, write a professional email to the supplier to get these products every week: \n\n- Wine (x10)\n- Eggs (x24)\n- Bread (x12)"
       }, {
         "title": "Code a snake game",
@@ -124,8 +124,40 @@ MODELS=`[
   }
 ]`
 
-OLD_MODELS=`[{"name":"bigcode/starcoder"}, {"name":"OpenAssistant/oasst-sft-6-llama-30b-xor"}, {"name":"mistralai/Mistral-7B-Instruct-v0.1"}]`
-TASK_MODEL='HuggingFaceH4/zephyr-7b-alpha'
+OLD_MODELS=`[{"name":"bigcode/starcoder"}, {"name":"OpenAssistant/oasst-sft-6-llama-30b-xor"}, {"name":"HuggingFaceH4/zephyr-7b-alpha"}]`
+
+TASK_MODEL='
+{
+  "name": "HuggingFaceH4/zephyr-7b-alpha",
+  "displayName": "HuggingFaceH4/zephyr-7b-alpha",
+  "description": "Zephyr 7B α is a fine-tune of Mistral 7B, released by the Hugging Face H4 RLHF team.",
+  "websiteUrl": "https://huggingface.co/HuggingFaceH4/zephyr-7b-alpha/",
+  "preprompt": "",
+  "chatPromptTemplate" : "<|system|>\n{{preprompt}}</s>\n{{#each messages}}{{#ifUser}}<|user|>\n{{content}}</s>\n<|assistant|>\n{{/ifUser}}{{#ifAssistant}}{{content}}</s>\n{{/ifAssistant}}{{/each}}",
+  "parameters": {
+    "temperature": 0.7,
+    "top_p": 0.95,
+    "repetition_penalty": 1.2,
+    "top_k": 50,
+    "truncate": 1000,
+    "max_new_tokens": 2048,
+    "stop": ["</s>", "<|>"]
+  },
+  "promptExamples": [
+    {
+      "title": "Write an email from bullet list",
+      "prompt": "As a restaurant owner, write a professional email to the supplier to get these products every week: \n\n- Wine (x10)\n- Eggs (x24)\n- Bread (x12)"
+    }, {
+      "title": "Code a snake game",
+      "prompt": "Code a basic snake game in python, give explanations for each step."
+    }, {
+      "title": "Assist in a task",
+      "prompt": "How do I make a delicious lemon cheesecake?"
+    }
+  ]
+}
+'
+
 
 APP_BASE="/chat"
 PUBLIC_ORIGIN=https://huggingface.co

src/lib/server/models.ts

@@ -37,70 +37,68 @@ const combinedEndpoint = endpoint.transform((data) => {
 	}
 });
 
-const modelsRaw = z
-	.array(
-		z.object({
-			/** Used as an identifier in DB */
-			id: z.string().optional(),
-			/** Used to link to the model page, and for inference */
-			name: z.string().min(1),
-			displayName: z.string().min(1).optional(),
-			description: z.string().min(1).optional(),
-			websiteUrl: z.string().url().optional(),
-			modelUrl: z.string().url().optional(),
-			datasetName: z.string().min(1).optional(),
-			datasetUrl: z.string().url().optional(),
-			userMessageToken: z.string().default(""),
-			userMessageEndToken: z.string().default(""),
-			assistantMessageToken: z.string().default(""),
-			assistantMessageEndToken: z.string().default(""),
-			messageEndToken: z.string().default(""),
-			preprompt: z.string().default(""),
-			prepromptUrl: z.string().url().optional(),
-			chatPromptTemplate: z
-				.string()
-				.default(
-					"{{preprompt}}" +
-						"{{#each messages}}" +
-						"{{#ifUser}}{{@root.userMessageToken}}{{content}}{{@root.userMessageEndToken}}{{/ifUser}}" +
-						"{{#ifAssistant}}{{@root.assistantMessageToken}}{{content}}{{@root.assistantMessageEndToken}}{{/ifAssistant}}" +
-						"{{/each}}" +
-						"{{assistantMessageToken}}"
-				),
-			promptExamples: z
-				.array(
-					z.object({
-						title: z.string().min(1),
-						prompt: z.string().min(1),
-					})
-				)
-				.optional(),
-			endpoints: z.array(combinedEndpoint).optional(),
-			parameters: z
-				.object({
-					temperature: z.number().min(0).max(1),
-					truncate: z.number().int().positive(),
-					max_new_tokens: z.number().int().positive(),
-					stop: z.array(z.string()).optional(),
-				})
-				.passthrough()
-				.optional(),
+const modelConfig = z.object({
+	/** Used as an identifier in DB */
+	id: z.string().optional(),
+	/** Used to link to the model page, and for inference */
+	name: z.string().min(1),
+	displayName: z.string().min(1).optional(),
+	description: z.string().min(1).optional(),
+	websiteUrl: z.string().url().optional(),
+	modelUrl: z.string().url().optional(),
+	datasetName: z.string().min(1).optional(),
+	datasetUrl: z.string().url().optional(),
+	userMessageToken: z.string().default(""),
+	userMessageEndToken: z.string().default(""),
+	assistantMessageToken: z.string().default(""),
+	assistantMessageEndToken: z.string().default(""),
+	messageEndToken: z.string().default(""),
+	preprompt: z.string().default(""),
+	prepromptUrl: z.string().url().optional(),
+	chatPromptTemplate: z
+		.string()
+		.default(
+			"{{preprompt}}" +
+				"{{#each messages}}" +
+				"{{#ifUser}}{{@root.userMessageToken}}{{content}}{{@root.userMessageEndToken}}{{/ifUser}}" +
+				"{{#ifAssistant}}{{@root.assistantMessageToken}}{{content}}{{@root.assistantMessageEndToken}}{{/ifAssistant}}" +
+				"{{/each}}" +
+				"{{assistantMessageToken}}"
+		),
+	promptExamples: z
+		.array(
+			z.object({
+				title: z.string().min(1),
+				prompt: z.string().min(1),
+			})
+		)
+		.optional(),
+	endpoints: z.array(combinedEndpoint).optional(),
+	parameters: z
+		.object({
+			temperature: z.number().min(0).max(1),
+			truncate: z.number().int().positive(),
+			max_new_tokens: z.number().int().positive(),
+			stop: z.array(z.string()).optional(),
 		})
-	)
-	.parse(JSON.parse(MODELS));
+		.passthrough()
+		.optional(),
+});
 
-export const models = await Promise.all(
-	modelsRaw.map(async (m) => ({
-		...m,
-		userMessageEndToken: m?.userMessageEndToken || m?.messageEndToken,
-		assistantMessageEndToken: m?.assistantMessageEndToken || m?.messageEndToken,
-		chatPromptRender: compileTemplate<ChatTemplateInput>(m.chatPromptTemplate, m),
-		id: m.id || m.name,
-		displayName: m.displayName || m.name,
-		preprompt: m.prepromptUrl ? await fetch(m.prepromptUrl).then((r) => r.text()) : m.preprompt,
-		parameters: { ...m.parameters, stop_sequences: m.parameters?.stop },
-	}))
-);
+const modelsRaw = z.array(modelConfig).parse(JSON.parse(MODELS));
+
+const processModel = async (m: z.infer<typeof modelConfig>) => ({
+	...m,
+	userMessageEndToken: m?.userMessageEndToken || m?.messageEndToken,
+	assistantMessageEndToken: m?.assistantMessageEndToken || m?.messageEndToken,
+	chatPromptRender: compileTemplate<ChatTemplateInput>(m.chatPromptTemplate, m),
+	id: m.id || m.name,
+	displayName: m.displayName || m.name,
+	preprompt: m.prepromptUrl ? await fetch(m.prepromptUrl).then((r) => r.text()) : m.preprompt,
+	parameters: { ...m.parameters, stop_sequences: m.parameters?.stop },
+});
+
+export const models = await Promise.all(modelsRaw.map(processModel));
 
 // Models that have been deprecated
 export const oldModels = OLD_MODELS
@@ -116,14 +114,18 @@ export const oldModels = OLD_MODELS
 			.map((m) => ({ ...m, id: m.id || m.name, displayName: m.displayName || m.name }))
 	: [];
 
-export type BackendModel = Optional<(typeof models)[0], "preprompt" | "parameters">;
-export type Endpoint = z.infer<typeof endpoint>;
-
 export const defaultModel = models[0];
 
-export const smallModel = models.find((m) => m.name === TASK_MODEL) || defaultModel;
-
 export const validateModel = (_models: BackendModel[]) => {
 	// Zod enum function requires 2 parameters
 	return z.enum([_models[0].id, ..._models.slice(1).map((m) => m.id)]);
 };
+
+// if `TASK_MODEL` is the name of a model we use it, else we try to parse `TASK_MODEL` as a model config itself
+export const smallModel = TASK_MODEL
+	? models.find((m) => m.name === TASK_MODEL) ||
+	  (await processModel(modelConfig.parse(JSON.parse(TASK_MODEL))))
+	: defaultModel;
+
+export type BackendModel = Optional<(typeof models)[0], "preprompt" | "parameters">;
+export type Endpoint = z.infer<typeof endpoint>;