Merge pull request #74 from replicate/mattt/prompt-template

Use prompt template to format messages

Author: mattt

app/api/route.js

@@ -1,6 +1,5 @@
 import Replicate from "replicate";
 import { ReplicateStream, StreamingTextResponse } from "ai";
-
 export const runtime = "edge";
 
 const replicate = new Replicate({
@@ -43,13 +42,12 @@ async function runLlama({
 }) {
   console.log("running llama");
 
-  const [owner, name] = model.split("/");
-
-  return await replicate.models.predictions.create(owner, name, {
+  return await replicate.predictions.create({
+    model: model,
     stream: true,
     input: {
       prompt: `${prompt}`,
-      system_prompt: systemPrompt,
+      prompt_template: "{prompt}",
       max_new_tokens: maxTokens,
       temperature: temperature,
       repetition_penalty: 1,

app/page.js

@@ -8,6 +8,8 @@ import EmptyState from "./components/EmptyState";
 import { Cog6ToothIcon, CodeBracketIcon } from "@heroicons/react/20/solid";
 import { useCompletion } from "ai/react";
 import { Toaster, toast } from "react-hot-toast";
+import { LlamaTemplate } from "../src/prompt_template";
+
 import { countTokens } from "./src/tokenizer.js";
 
 const MODELS = [
@@ -38,6 +40,20 @@ const MODELS = [
   },
 ];
 
+const llamaTemplate = LlamaTemplate();
+
+const generatePrompt = (template, systemPrompt, messages) => {
+  const chat = messages.map((message) => ({
+    "role": message.isUser ? "user" : "assistant",
+    "content": message.text,
+  }));
+
+  return template([{
+    "role": "system",
+    "content": systemPrompt,
+  }, ...chat]);
+};
+
 function CTA({ shortenedModelName }) {
   if (shortenedModelName == "Llava") {
     return (
@@ -141,7 +157,6 @@ export default function HomePage() {
 
   const handleFileUpload = (file) => {
     if (file) {
-      console.log(file);
       // determine if file is image or audio
       if (
         ["audio/mpeg", "audio/wav", "audio/ogg"].includes(
@@ -192,16 +207,8 @@ export default function HomePage() {
       isUser: true,
     });
 
-    const generatePrompt = (messages) => {
-      return messages
-        .map((message) =>
-          message.isUser ? `[INST] ${message.text} [/INST]` : `${message.text}`
-        )
-        .join("\n");
-    };
-
     // Generate initial prompt and calculate tokens
-    let prompt = `${generatePrompt(messageHistory)}\n`;
+    let prompt = `${generatePrompt(llamaTemplate, systemPrompt, messageHistory)}\n`;
     // Check if we exceed max tokens and truncate the message history if so.
     while (countTokens(prompt) > MAX_TOKENS) {
       if (messageHistory.length < 3) {
@@ -216,7 +223,7 @@ export default function HomePage() {
       messageHistory.splice(1, 2);
 
       // Recreate the prompt
-      prompt = `${SNIP}\n${generatePrompt(messageHistory)}\n`;
+      prompt = `${SNIP}\n${generatePrompt(llamaTemplate, systemPrompt, messageHistory)}\n`;
     }
 
     setMessages(messageHistory);