85 lines
3.2 KiB
TypeScript
85 lines
3.2 KiB
TypeScript
import type { Message } from '../shared/types';
|
|
import {
|
|
experimental_buildLlama2Prompt,
|
|
experimental_buildOpenAssistantPrompt,
|
|
experimental_buildStarChatBetaPrompt,
|
|
} from './huggingface';
|
|
|
|
describe('buildStarChatBetaPrompt', () => {
|
|
it('should return a string with user, assistant, and system messages', () => {
|
|
const messages: Pick<Message, 'content' | 'role'>[] = [
|
|
{ content: 'You are a chat bot.', role: 'system' },
|
|
{ content: 'Hello!', role: 'user' },
|
|
{ content: 'Hi there!', role: 'assistant' },
|
|
];
|
|
|
|
const expectedPrompt = `<|system|>\nYou are a chat bot.<|end|>\n<|user|>\nHello!<|end|>\n<|assistant|>\nHi there!<|end|>\n<|assistant|>`;
|
|
const prompt = experimental_buildStarChatBetaPrompt(messages);
|
|
expect(prompt).toEqual(expectedPrompt);
|
|
});
|
|
|
|
it('should throw an error if a function message is included', () => {
|
|
const messages: Pick<Message, 'content' | 'role'>[] = [
|
|
{ content: 'someFunction()', role: 'function' },
|
|
];
|
|
expect(() => experimental_buildStarChatBetaPrompt(messages)).toThrow();
|
|
});
|
|
});
|
|
|
|
describe('buildOpenAssistantPrompt', () => {
|
|
it('should return a string with user and assistant messages', () => {
|
|
const messages: Pick<Message, 'content' | 'role'>[] = [
|
|
{ content: 'Hello!', role: 'user' },
|
|
{ content: 'Hi there!', role: 'assistant' },
|
|
];
|
|
|
|
const expectedPrompt =
|
|
'<|prompter|>Hello!<|endoftext|><|assistant|>Hi there!<|endoftext|><|assistant|>';
|
|
const prompt = experimental_buildOpenAssistantPrompt(messages);
|
|
expect(prompt).toEqual(expectedPrompt);
|
|
});
|
|
|
|
it('should throw an error if a function message is included', () => {
|
|
const messages: Pick<Message, 'content' | 'role'>[] = [
|
|
{ content: 'someFunction()', role: 'function' },
|
|
];
|
|
expect(() => experimental_buildOpenAssistantPrompt(messages)).toThrow();
|
|
});
|
|
});
|
|
|
|
describe('buildLlamaPrompt', () => {
|
|
it('should return a string with user instruction', () => {
|
|
const messages: Pick<Message, 'content' | 'role'>[] = [
|
|
{ content: 'Hello, how are you?', role: 'user' },
|
|
];
|
|
|
|
const expectedPrompt = '<s>[INST] Hello, how are you? [/INST]';
|
|
const prompt = experimental_buildLlama2Prompt(messages);
|
|
expect(prompt).toEqual(expectedPrompt);
|
|
});
|
|
|
|
it('should return a string with system, user and assistant messages', () => {
|
|
const messages: Pick<Message, 'content' | 'role'>[] = [
|
|
{
|
|
content: 'You are helpful assistant, but you are drunk, hick',
|
|
role: 'system',
|
|
},
|
|
{ content: 'Hi there!', role: 'user' },
|
|
{ content: 'Sup, partner!', role: 'assistant' },
|
|
{ content: 'What are you doing?', role: 'user' },
|
|
];
|
|
|
|
const expectedPrompt =
|
|
'<s>[INST] <<SYS>>\nYou are helpful assistant, but you are drunk, hick\n<</SYS>>\n\nHi there! [/INST] Sup, partner!</s><s>[INST] What are you doing? [/INST]';
|
|
const prompt = experimental_buildLlama2Prompt(messages);
|
|
expect(prompt).toEqual(expectedPrompt);
|
|
});
|
|
|
|
it('should throw an error if a function message is included', () => {
|
|
const messages: Pick<Message, 'content' | 'role'>[] = [
|
|
{ content: 'someFunction()', role: 'function' },
|
|
];
|
|
expect(() => experimental_buildLlama2Prompt(messages)).toThrow();
|
|
});
|
|
});
|