replicate/llama2 provider

Still need to fix the types but it runs
This commit is contained in:
Kyle Corbitt
2023-07-20 19:55:03 -07:00
parent 332a2101c0
commit 847753c32b
14 changed files with 202 additions and 29 deletions

View File

@@ -0,0 +1,74 @@
import { type ModelProvider } from "../types";
import { getCompletion } from "./getCompletion";
const supportedModels = ["7b-chat", "13b-chat", "70b-chat"] as const;
type SupportedModel = (typeof supportedModels)[number];
export type ReplicateLlama2Input = {
model: SupportedModel;
prompt: string;
stream?: boolean;
max_length?: number;
temperature?: number;
top_p?: number;
repetition_penalty?: number;
debug?: boolean;
};
export type ReplicateLlama2Output = string[];
export type ReplicateLlama2Provider = ModelProvider<
SupportedModel,
ReplicateLlama2Input,
ReplicateLlama2Output
>;
const modelProvider: ReplicateLlama2Provider = {
name: "OpenAI ChatCompletion",
models: {
"7b-chat": {},
"13b-chat": {},
"70b-chat": {},
},
getModel: (input) => {
if (supportedModels.includes(input.model)) return input.model;
return null;
},
inputSchema: {
type: "object",
properties: {
model: {
type: "string",
enum: supportedModels as unknown as string[],
},
prompt: {
type: "string",
},
stream: {
type: "boolean",
},
max_length: {
type: "number",
},
temperature: {
type: "number",
},
top_p: {
type: "number",
},
repetition_penalty: {
type: "number",
},
debug: {
type: "boolean",
},
},
required: ["model", "prompt"],
},
shouldStream: (input) => input.stream ?? false,
getCompletion,
};
export default modelProvider;