♻️ Simplify setup, prepare for multi-models (#156)

.env

@@ -1,27 +1,29 @@
-# Use .env.local to change these variables, or directly change your env
+# Use .env.local to change these variables
 # DO NOT EDIT THIS FILE WITH SENSITIVE DATA
 
 MONGODB_URL=#your mongodb URL here
 MONGODB_DB_NAME=chat-ui
 COOKIE_NAME=hf-chat
+HF_ACCESS_TOKEN=#hf_<token> from from https://huggingface.co/settings/token
+MODELS=`["OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5"]`
+# Alternative syntax (all fields are optional except 'name'):
+# MODELS=`[{
+#  "name": "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5", 
+#  "displayName": "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
+#  "endpoints": [{
+#    "url": "https://api-inference.huggingface.co/models/OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
+#    "authorization": "Bearer hf_<token>",
+#    "weight": 1
+#  }]
+# }]
 
 # Increase depending on the model
 PUBLIC_MAX_INPUT_TOKENS=1000
 PUBLIC_ORIGIN=#https://hf.co
-PUBLIC_MODEL_NAME=OpenAssistant/oasst-sft-6-llama-30b # public facing link
-PUBLIC_MODEL_ID=OpenAssistant/oasst-sft-6-llama-30b-xor # used to link to model page
 PUBLIC_DISABLE_INTRO_TILES=false
 PUBLIC_USER_MESSAGE_TOKEN=<|prompter|>
 PUBLIC_ASSISTANT_MESSAGE_TOKEN=<|assistant|>
 PUBLIC_SEP_TOKEN=</s>
 PUBLIC_PREPROMPT="Below are a series of dialogues between various people and an AI assistant. The AI tries to be helpful, polite, honest, sophisticated, emotionally aware, and humble-but-knowledgeable. The assistant is happy to help with almost anything, and will do its best to understand exactly what is needed. It also tries to avoid giving false or misleading information, and it caveats when it isn't entirely sure about the right answer. That said, the assistant is practical and really does its best, and doesn't let caution get too much in the way of being useful."
 PUBLIC_GOOGLE_ANALYTICS_ID=#G-XXXXXXXX / Leave empty to disable
-PUBLIC_DEPRECATED_GOOGLE_ANALYTICS_ID=#UA-XXXXXXXX-X / Leave empty to disable
-
-# Copy this in .env.local with and replace "hf_<token>" your HF token from https://huggingface.co/settings/token
-# You can also change the model from OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5 to your own model
-MODEL_ENDPOINTS=`[{
-  "endpoint": "https://api-inference.huggingface.co/models/OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5", 
-  "authorization": "Bearer hf_<token>", 
-  "weight": 1
-}]`
+PUBLIC_DEPRECATED_GOOGLE_ANALYTICS_ID=#UA-XXXXXXXX-X / Leave empty to disable

README.md

@@ -31,30 +31,32 @@ Basically you need to create a `.env.local` with the following contents:
 
 ```
 MONGODB_URL=<url to mongo, for example a free MongoDB Atlas sandbox instance>
-MODEL_ENDPOINTS=`[{
-  "endpoint": "https://api-inference.huggingface.co/models/OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
-  "authorization": "Bearer <hf_token>",
-  "weight": 1
-}]`
+HF_ACCESS_TOKEN=<your HF access token from https://huggingface.co/settings/tokens>
 ```
 
-Where the contents in `<...>` are replaced by the MongoDB URL and your [HF Access Token](https://huggingface.co/settings/tokens).
-
 ## Duplicating to a Space
 
 Create a `DOTENV_LOCAL` secret to your space with the following contents:
 
 ```
 MONGODB_URL=<url to mongo, for example a free MongoDB Atlas sandbox instance>
-MODEL_ENDPOINTS=`[{
-  "endpoint": "https://api-inference.huggingface.co/models/OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
-  "authorization": "Bearer <hf_token>",
-  "weight": 1
-}]`
+HF_ACCESS_TOKEN=<your HF access token from https://huggingface.co/settings/tokens>
 ```
 
 Where the contents in `<...>` are replaced by the MongoDB URL and your [HF Access Token](https://huggingface.co/settings/tokens).
 
+## Running Local Inference
+
+Both the example above use the HF Inference API or HF Endpoints API.
+
+If you want to run the model locally, you need to run this inference server locally: https://github.com/huggingface/text-generation-inference
+
+And add this to your `.env.local`:
+
+```
+MODELS=`[{"name": "...", "endpoints": [{"url": "127.0.0.1:8080/generate_stream"}]}]`
+```
+
 ## Building
 
 To create a production version of your app:

src/hooks.server.ts

@@ -6,6 +6,7 @@ import {
 	PUBLIC_DEPRECATED_GOOGLE_ANALYTICS_ID,
 } from "$env/static/public";
 import { addYears } from "date-fns";
+import { inspect } from "node:util";
 
 export const handle: Handle = async ({ event, resolve }) => {
 	const token = event.cookies.get(COOKIE_NAME);

src/lib/components/chat/ChatIntroduction.svelte

@@ -1,15 +1,11 @@
 <script lang="ts">
-	import {
-		PUBLIC_DISABLE_INTRO_TILES,
-		PUBLIC_MODEL_ID,
-		PUBLIC_MODEL_NAME,
-		PUBLIC_VERSION,
-	} from "$env/static/public";
-
+	import { PUBLIC_DISABLE_INTRO_TILES, PUBLIC_VERSION } from "$env/static/public";
 	import Logo from "$lib/components/icons/Logo.svelte";
 	import CarbonArrowUpRight from "~icons/carbon/arrow-up-right";
 	import CarbonEarth from "~icons/carbon/earth";
 	import { createEventDispatcher } from "svelte";
+
+	export let currentModel: { name: string; displayName: string };
 	const dispatch = createEventDispatcher<{ message: string }>();
 </script>
 
@@ -36,13 +32,13 @@
 		<div class="overflow-hidden rounded-xl border dark:border-gray-800">
 			<div class="p-3">
 				<div class="text-sm text-gray-600 dark:text-gray-400">Current Model</div>
-				<div class="font-semibold">{PUBLIC_MODEL_NAME}</div>
+				<div class="font-semibold">{currentModel.displayName}</div>
 			</div>
 			<div
 				class="flex items-center gap-5 rounded-xl bg-gray-100 px-3 py-2 text-sm text-gray-600 dark:bg-gray-800 dark:text-gray-300"
 			>
 				<a
-					href="https://huggingface.co/{PUBLIC_MODEL_ID}"
+					href="https://huggingface.co/{currentModel.name}"
 					target="_blank"
 					rel="noreferrer"
 					class="flex items-center hover:underline"

src/lib/components/chat/ChatMessages.svelte

@@ -13,6 +13,7 @@
 	export let messages: Message[];
 	export let loading: boolean;
 	export let pending: boolean;
+	export let currentModel: { name: string; displayName: string };
 
 	let chatContainer: HTMLElement;
 
@@ -40,7 +41,7 @@
 				on:retry={() => dispatch("retry", { id: message.id, content: message.content })}
 			/>
 		{:else}
-			<ChatIntroduction on:message />
+			<ChatIntroduction on:message {currentModel} />
 		{/each}
 		{#if pending}
 			<ChatMessage message={{ from: "assistant", content: "", id: randomUUID() }} />

src/lib/components/chat/ChatWindow.svelte

@@ -8,12 +8,12 @@
 	import ChatMessages from "./ChatMessages.svelte";
 	import ChatInput from "./ChatInput.svelte";
 	import StopGeneratingBtn from "../StopGeneratingBtn.svelte";
-	import { PUBLIC_MODEL_ID, PUBLIC_MODEL_NAME } from "$env/static/public";
 
 	export let messages: Message[] = [];
 	export let disabled = false;
 	export let loading = false;
 	export let pending = false;
+	export let currentModel: { name: string; displayName: string };
 
 	let message: string;
 
@@ -35,6 +35,7 @@
 	<ChatMessages
 		{loading}
 		{pending}
+		{currentModel}
 		{messages}
 		on:message
 		on:retry={(ev) => {
@@ -73,10 +74,10 @@
 		<div class="mt-2 flex justify-between self-stretch px-1 text-xs text-gray-400/90 max-sm:gap-2">
 			<p>
 				Model: <a
-					href="https://huggingface.co/{PUBLIC_MODEL_ID}"
+					href="https://huggingface.co/{currentModel.name}"
 					target="_blank"
 					rel="noreferrer"
-					class="hover:underline">{PUBLIC_MODEL_NAME}</a
+					class="hover:underline">{currentModel.displayName}</a
 				> <span class="max-sm:hidden">·</span><br class="sm:hidden" /> Generated content may be inaccurate
 				or false.
 			</p>

src/lib/server/modelEndpoint.ts

@@ -1,14 +1,38 @@
-import { MODEL_ENDPOINTS } from "$env/static/private";
+import { HF_ACCESS_TOKEN } from "$env/static/private";
 import { sum } from "$lib/utils/sum";
-
-const endpoints: Array<{ endpoint: string; authorization: string; weight: number }> =
-	JSON.parse(MODEL_ENDPOINTS);
-const totalWeight = sum(endpoints.map((e) => e.weight));
+import { models } from "./models";
 
 /**
  * Find a random load-balanced endpoint
  */
-export function modelEndpoint(): { endpoint: string; authorization: string; weight: number } {
+export function modelEndpoint(model: string): {
+	url: string;
+	authorization: string;
+	weight: number;
+} {
+	const modelDefinition = models.find(
+		(m) => m === model || (typeof m === "object" && m.name === model)
+	);
+	if (!modelDefinition) {
+		throw new Error(`Invalid model: ${model}`);
+	}
+	if (typeof modelDefinition === "string") {
+		return {
+			url: `https://api-inference.huggingface.co/models/${modelDefinition}`,
+			authorization: `Bearer ${HF_ACCESS_TOKEN}`,
+			weight: 1,
+		};
+	}
+	if (!modelDefinition.endpoints) {
+		return {
+			url: `https://api-inference.huggingface.co/models/${modelDefinition.name}`,
+			authorization: `Bearer ${HF_ACCESS_TOKEN}`,
+			weight: 1,
+		};
+	}
+	const endpoints = modelDefinition.endpoints;
+	const totalWeight = sum(endpoints.map((e) => e.weight));
+
 	let random = Math.random() * totalWeight;
 	for (const endpoint of endpoints) {
 		if (random < endpoint.weight) {

src/lib/server/models.ts

@@ -0,0 +1,30 @@
+import { HF_ACCESS_TOKEN, MODELS } from "$env/static/private";
+import { z } from "zod";
+
+export const models = z
+	.array(
+		z.union([
+			z.string().min(1),
+			z.object({
+				name: z.string().min(1),
+				displayName: z.string().min(1).optional(),
+				endpoints: z
+					.array(
+						z.object({
+							url: z.string().url(),
+							authorization: z.string().min(1).default(`Bearer ${HF_ACCESS_TOKEN}`),
+							weight: z.number().int().positive().default(1),
+						})
+					)
+					.optional(),
+			}),
+		])
+	)
+	.parse(JSON.parse(MODELS));
+
+export const modelNames: Array<{ name: string; displayName: string }> = models.map((m) =>
+	typeof m === "string"
+		? { name: m, displayName: m }
+		: { name: m.name, displayName: m.displayName ?? m.name }
+);
+export const defaultModel = modelNames[0];

src/lib/types/Message.ts

@@ -2,4 +2,6 @@ export interface Message {
 	from: "user" | "assistant";
 	id: ReturnType<typeof crypto.randomUUID>;
 	content: string;
+	// Only for "assistant" messages
+	model?: string;
 }

src/routes/+layout.server.ts

@@ -2,6 +2,7 @@ import type { LayoutServerLoad } from "./$types";
 import { collections } from "$lib/server/database";
 import type { Conversation } from "$lib/types/Conversation";
 import { UrlDependency } from "$lib/types/UrlDependency";
+import { modelNames } from "$lib/server/models";
 
 export const load: LayoutServerLoad = async ({ locals, depends }) => {
 	const { conversations } = collections;
@@ -29,5 +30,6 @@ export const load: LayoutServerLoad = async ({ locals, depends }) => {
 			shareConversationsWithModelAuthors: settings?.shareConversationsWithModelAuthors ?? true,
 			ethicsModalAcceptedAt: settings?.ethicsModalAcceptedAt ?? null,
 		},
+		models: modelNames,
 	};
 };

src/routes/+page.svelte

@@ -5,6 +5,7 @@
 	import { ERROR_MESSAGES, error } from "$lib/stores/errors";
 	import { pendingMessage } from "$lib/stores/pendingMessage";
 
+	export let data;
 	let loading = false;
 
 	async function createConversation(message: string) {
@@ -39,4 +40,8 @@
 	}
 </script>
 
-<ChatWindow on:message={(ev) => createConversation(ev.detail)} {loading} />
+<ChatWindow
+	on:message={(ev) => createConversation(ev.detail)}
+	{loading}
+	currentModel={data.models[0]}
+/>

src/routes/conversation/[id]/+page.svelte

@@ -181,4 +181,5 @@
 	on:retry={(message) => writeMessage(message.detail.content, message.detail.id)}
 	on:share={() => shareConversation($page.params.id, data.title)}
 	on:stop={() => (isAborted = true)}
+	currentModel={data.models[0]}
 />

src/routes/conversation/[id]/+server.ts

@@ -3,6 +3,7 @@ import { buildPrompt } from "$lib/buildPrompt.js";
 import { abortedGenerations } from "$lib/server/abortedGenerations.js";
 import { collections } from "$lib/server/database.js";
 import { modelEndpoint } from "$lib/server/modelEndpoint.js";
+import { defaultModel, modelNames } from "$lib/server/models.js";
 import type { Message } from "$lib/types/Message.js";
 import { concatUint8Arrays } from "$lib/utils/concatUint8Arrays.js";
 import { streamToAsyncIterable } from "$lib/utils/streamToAsyncIterable";
@@ -30,10 +31,14 @@ export async function POST({ request, fetch, locals, params }) {
 	const json = await request.json();
 	const {
 		inputs: newPrompt,
+		model,
 		options: { id: messageId, is_retry },
 	} = z
 		.object({
 			inputs: z.string().trim().min(1),
+			model: z
+				.enum([modelNames[0].name, ...modelNames.slice(1).map((m) => m.name)])
+				.default(defaultModel.name),
 			options: z.object({
 				id: z.optional(z.string().uuid()),
 				is_retry: z.optional(z.boolean()),
@@ -66,11 +71,11 @@ export async function POST({ request, fetch, locals, params }) {
 	}
 	const prompt = buildPrompt(messages);
 
-	const randomEndpoint = modelEndpoint();
+	const randomEndpoint = modelEndpoint(model);
 
 	const abortController = new AbortController();
 
-	const resp = await fetch(randomEndpoint.endpoint, {
+	const resp = await fetch(randomEndpoint.url, {
 		headers: {
 			"Content-Type": request.headers.get("Content-Type") ?? "application/json",
 			Authorization: randomEndpoint.authorization,
@@ -99,7 +104,7 @@ export async function POST({ request, fetch, locals, params }) {
 
 		generated_text = trimSuffix(trimPrefix(generated_text, "<|startoftext|>"), PUBLIC_SEP_TOKEN);
 
-		messages.push({ from: "assistant", content: generated_text, id: crypto.randomUUID() });
+		messages.push({ from: "assistant", content: generated_text, id: crypto.randomUUID(), model });
 
 		await collections.conversations.updateOne(
 			{

src/routes/conversation/[id]/summarize/+server.ts

@@ -2,6 +2,7 @@ import { PUBLIC_MAX_INPUT_TOKENS, PUBLIC_SEP_TOKEN } from "$env/static/public";
 import { buildPrompt } from "$lib/buildPrompt";
 import { collections } from "$lib/server/database.js";
 import { modelEndpoint } from "$lib/server/modelEndpoint.js";
+import { defaultModel } from "$lib/server/models.js";
 import { trimPrefix } from "$lib/utils/trimPrefix.js";
 import { trimSuffix } from "$lib/utils/trimSuffix.js";
 import { textGeneration } from "@huggingface/inference";
@@ -40,10 +41,10 @@ export async function POST({ params, locals, fetch }) {
 		return_full_text: false,
 	};
 
-	const endpoint = modelEndpoint();
+	const endpoint = modelEndpoint(defaultModel.name);
 	let { generated_text } = await textGeneration(
 		{
-			model: endpoint.endpoint,
+			model: endpoint.url,
 			inputs: prompt,
 			parameters,
 		},

src/routes/r/[id]/+page.svelte

@@ -71,5 +71,6 @@
 			})
 			.finally(() => (loading = false))}
 	messages={data.messages}
+	currentModel={data.models[0]}
 	{loading}
 />