@llumiverse/drivers 0.8.0

package/src/huggingface_ie.ts

@@ -0,0 +1,269 @@
+import {
+    AIModel,
+    AIModelStatus,
+    AbstractDriver,
+    BuiltinProviders,
+    DriverOptions,
+    ExecutionOptions,
+    PromptFormats
+} from "@llumiverse/core";
+import { transformAsyncIterator } from "@llumiverse/core/async";
+import {
+    HfInference,
+    HfInferenceEndpoint,
+    TextGenerationStreamOutput
+} from "@huggingface/inference";
+import { FetchClient } from "api-fetch-client";
+
+export interface HuggingFaceIEDriverOptions extends DriverOptions {
+    apiKey: string;
+    endpoint_url: string;
+}
+
+export class HuggingFaceIEDriver extends AbstractDriver<HuggingFaceIEDriverOptions, string> {
+    service: FetchClient;
+    provider = BuiltinProviders.huggingface_ie;
+    _executor?: HfInferenceEndpoint;
+    defaultFormat = PromptFormats.genericTextLLM;
+
+    constructor(
+        options: HuggingFaceIEDriverOptions
+    ) {
+        super(options);
+        if (!options.endpoint_url) {
+            throw new Error(`Endpoint URL is required for ${this.provider}`);
+        }
+        this.service = new FetchClient(this.options.endpoint_url);
+        this.service.headers["Authorization"] = `Bearer ${this.options.apiKey}`;
+    }
+
+    async getModelURLEndpoint(
+        modelId: string
+    ): Promise<{ url: string; status: string; }> {
+        const res = (await this.service.get(`/${modelId}`)) as HuggingFaceIEModel;
+        return {
+            url: res.status.url,
+            status: getStatus(res),
+        };
+    }
+
+    async getExecutor(model: string) {
+        if (!this._executor) {
+            const endpoint = await this.getModelURLEndpoint(model);
+            if (!endpoint.url)
+                throw new Error(
+                    `Endpoint URL not found for model ${model}`
+                );
+            if (endpoint.status !== AIModelStatus.Available)
+                throw new Error(
+                    `Endpoint ${model} is not running - current status: ${endpoint.status}`
+                );
+
+            this._executor = new HfInference(this.options.apiKey).endpoint(
+                endpoint.url
+            );
+        }
+        return this._executor;
+    }
+
+    async requestCompletionStream(prompt: string, options: ExecutionOptions) {
+        const executor = await this.getExecutor(options.model);
+        const req = executor.textGenerationStream({
+            inputs: prompt,
+            parameters: {
+                temperature: options.temperature,
+                max_new_tokens: options.max_tokens,
+            },
+        });
+
+        return transformAsyncIterator(req, (val: TextGenerationStreamOutput) => {
+            //special like <s> are not part of the result
+            if (val.token.special) return "";
+            return val.token.text;
+        });
+    }
+
+    async requestCompletion(prompt: string, options: ExecutionOptions) {
+        const executor = await this.getExecutor(options.model);
+        const res = await executor.textGeneration({
+            inputs: prompt,
+            parameters: {
+                temperature: options.temperature,
+                max_new_tokens: options.max_tokens,
+            },
+        });
+
+        return {
+            result: res.generated_text,
+            token_usage: {
+                result: res.generated_text.length,
+                prompt: prompt.length,
+                total: res.generated_text.length + prompt.length,
+            },
+        };
+
+    }
+
+    // ============== management API ==============
+
+    // Not implemented
+    async listTrainableModels(): Promise<AIModel<string>[]> {
+        return [];
+    }
+
+    async listModels(): Promise<AIModel[]> {
+        const res = await this.service.get("/");
+        const hfModels = res.items as HuggingFaceIEModel[];
+
+        const models: AIModel[] = hfModels.map((model: HuggingFaceIEModel) => ({
+            id: model.name,
+            name: `${model.name} [${model.model.repository}:${model.model.task}]`,
+            provider: this.provider,
+            tags: [model.model.task],
+            status: getStatus(model),
+        }));
+
+        return models;
+    }
+
+    async validateConnection(): Promise<boolean> {
+        try {
+            await this.service.get("/models");
+            return true;
+        } catch (error) {
+            return false;
+        }
+    }
+
+    async generateEmbeddings(content: string, model?: string): Promise<{ embeddings: number[], model: string; }> {
+        this.logger?.debug(`[Huggingface] Generating embeddings for ${content} on ${model}`);
+        throw new Error("Method not implemented.");
+    }
+
+}
+
+//get status from HF status
+function getStatus(hfModel: HuggingFaceIEModel): AIModelStatus {
+    //[ pending, initializing, updating, updateFailed, running, paused, failed, scaledToZero ]
+    switch (hfModel.status.state) {
+        case "running":
+            return AIModelStatus.Available;
+        case "initializing":
+            return AIModelStatus.Pending;
+        case "updating":
+            return AIModelStatus.Pending;
+        case "updateFailed":
+            return AIModelStatus.Unavailable;
+        case "paused":
+            return AIModelStatus.Stopped;
+        case "failed":
+            return AIModelStatus.Unavailable;
+        case "scaledToZero":
+            return AIModelStatus.Available;
+        default:
+            return AIModelStatus.Unknown;
+    }
+}
+
+interface HuggingFaceIEModel {
+    accountId: string;
+    compute: {
+        accelerator: string;
+        instanceSize: string;
+        instanceType: string;
+        scaling: {
+            maxReplica: number;
+            minReplica: number;
+        };
+    };
+    model: {
+        framework: string;
+        image: {
+            huggingface: {};
+        };
+        repository: string;
+        revision: string;
+        task: string;
+    };
+    name: string;
+    provider: {
+        region: string;
+        vendor: string;
+    };
+    status: {
+        createdAt: string;
+        createdBy: {
+            id: string;
+            name: string;
+        };
+        message: string;
+        private: {
+            serviceName: string;
+        };
+        readyReplica: number;
+        state: string;
+        targetReplica: number;
+        updatedAt: string;
+        updatedBy: {
+            id: string;
+            name: string;
+        };
+        url: string;
+    };
+    type: string;
+}
+
+/*
+Example of model returned by the API
+{
+    "items": [
+      {
+        "accountId": "string",
+        "compute": {
+          "accelerator": "cpu",
+          "instanceSize": "large",
+          "instanceType": "c6i",
+          "scaling": {
+            "maxReplica": 8,
+            "minReplica": 2
+          }
+        },
+        "model": {
+          "framework": "custom",
+          "image": {
+            "huggingface": {}
+          },
+          "repository": "gpt2",
+          "revision": "6c0e6080953db56375760c0471a8c5f2929baf11",
+          "task": "text-classification"
+        },
+        "name": "my-endpoint",
+        "provider": {
+          "region": "us-east-1",
+          "vendor": "aws"
+        },
+        "status": {
+          "createdAt": "2023-10-19T05:04:17.305Z",
+          "createdBy": {
+            "id": "string",
+            "name": "string"
+          },
+          "message": "Endpoint is ready",
+          "private": {
+            "serviceName": "string"
+          },
+          "readyReplica": 2,
+          "state": "pending",
+          "targetReplica": 4,
+          "updatedAt": "2023-10-19T05:04:17.305Z",
+          "updatedBy": {
+            "id": "string",
+            "name": "string"
+          },
+          "url": "https://endpoint-id.region.vendor.endpoints.huggingface.cloud"
+        },
+        "type": "public"
+      }
+    ]
+  }
+*/

package/src/index.ts

@@ -0,0 +1,7 @@
+export * from "./bedrock/index.js";
+export * from "./huggingface_ie.js";
+export * from "./openai.js";
+export * from "./replicate.js";
+export * from "./vertexai/index.js";
+export * from "./togetherai/index.js";
+export * from "./test/index.js";

package/src/openai.ts

@@ -0,0 +1,254 @@
+import {
+    AIModel,
+    AbstractDriver,
+    BuiltinProviders,
+    Completion,
+    DataSource,
+    DriverOptions,
+    ExecutionOptions,
+    ExecutionTokenUsage,
+    ModelType,
+    PromptFormats,
+    PromptSegment,
+    TrainingJob,
+    TrainingJobStatus,
+    TrainingOptions,
+    TrainingPromptOptions
+} from "@llumiverse/core";
+import { asyncMap } from "@llumiverse/core/async";
+import OpenAI from "openai";
+import { Stream } from "openai/streaming";
+
+const supportFineTunning = new Set([
+    "gpt-3.5-turbo-1106",
+    "gpt-3.5-turbo-0613",
+    "babbage-002",
+    "davinci-002",
+    "gpt-4-0613"
+]);
+
+export interface OpenAIDriverOptions extends DriverOptions {
+    apiKey: string;
+}
+
+export class OpenAIDriver extends AbstractDriver<
+    OpenAIDriverOptions,
+    OpenAI.Chat.Completions.ChatCompletionMessageParam[]
+> {
+    inputContentTypes: string[] = ["text/plain"];
+    generatedContentTypes: string[] = ["text/plain"];
+    service: OpenAI;
+    provider = BuiltinProviders.openai;
+    defaultFormat = PromptFormats.openai;
+
+    constructor(opts: OpenAIDriverOptions) {
+        super(opts);
+        this.service = new OpenAI({
+            apiKey: opts.apiKey,
+        });
+    }
+
+    createPrompt(segments: PromptSegment[], opts: ExecutionOptions): OpenAI.Chat.Completions.ChatCompletionMessageParam[] {
+        // openai only supports opanai format -  force the format
+        return super.createPrompt(segments, { ...opts, format: PromptFormats.openai })
+    }
+
+    extractDataFromResponse(
+        options: ExecutionOptions,
+        result: OpenAI.Chat.Completions.ChatCompletion
+    ): Completion {
+        const tokenInfo: ExecutionTokenUsage = {
+            prompt: result.usage?.prompt_tokens,
+            result: result.usage?.completion_tokens,
+            total: result.usage?.total_tokens,
+        };
+
+        //if no schema, return content
+        if (!options.resultSchema) {
+            return {
+                result: result.choices[0]?.message.content as string,
+                token_usage: tokenInfo,
+            }
+        }
+
+        //we have a schema: get the content and return after validation
+        const data = result.choices[0]?.message.function_call?.arguments as any;
+        if (!data) {
+            this.logger?.error("[OpenAI] Response is not valid", result);
+            throw new Error("Response is not valid: no data");
+        }
+
+        return {
+            result: data,
+            token_usage: tokenInfo
+        };
+    }
+
+    async requestCompletionStream(prompt: OpenAI.Chat.Completions.ChatCompletionMessageParam[], options: ExecutionOptions): Promise<any> {
+        const mapFn = options.resultSchema
+            ? (chunk: OpenAI.Chat.Completions.ChatCompletionChunk) => {
+                return (
+                    chunk.choices[0]?.delta?.function_call?.arguments ?? ""
+                );
+            }
+            : (chunk: OpenAI.Chat.Completions.ChatCompletionChunk) => {
+                return chunk.choices[0]?.delta?.content ?? "";
+            };
+
+        const stream = (await this.service.chat.completions.create({
+            stream: true,
+            model: options.model,
+            messages: prompt,
+            temperature: options.temperature,
+            n: 1,
+            max_tokens: options.max_tokens,
+            functions: options.resultSchema
+                ? [
+                    {
+                        name: "format_output",
+                        parameters: options.resultSchema as any,
+                    },
+                ]
+                : undefined,
+            function_call: options.resultSchema
+                ? { name: "format_output" }
+                : undefined,
+        })) as Stream<OpenAI.Chat.Completions.ChatCompletionChunk>;
+
+        return asyncMap(stream, mapFn);
+    }
+
+    async requestCompletion(prompt: OpenAI.Chat.Completions.ChatCompletionMessageParam[], options: ExecutionOptions): Promise<any> {
+        const functions = options.resultSchema
+            ? [
+                {
+                    name: "format_output",
+                    parameters: options.resultSchema as any,
+                },
+            ]
+            : undefined;
+
+        const res = await this.service.chat.completions.create({
+            stream: false,
+            model: options.model,
+            messages: prompt,
+            temperature: options.temperature,
+            n: 1,
+            max_tokens: options.max_tokens,
+            functions: functions,
+            function_call: options.resultSchema
+                ? { name: "format_output" }
+                : undefined,
+        });
+
+        return this.extractDataFromResponse(options, res);
+    }
+
+    createTrainingPrompt(options: TrainingPromptOptions): string {
+        if (options.model.includes("gpt")) {
+            return super.createTrainingPrompt(options);
+        } else {
+            // babbage, davinci not yet implemented 
+            throw new Error("Unsupported model for training: " + options.model);
+        }
+    }
+
+    async startTraining(dataset: DataSource, options: TrainingOptions): Promise<TrainingJob> {
+        const url = await dataset.getURL();
+        const file = await this.service.files.create({
+            file: await fetch(url),
+            purpose: "fine-tune",
+        });
+
+        const job = await this.service.fineTuning.jobs.create({
+            training_file: file.id,
+            model: options.model,
+            hyperparameters: options.params
+        })
+
+        return jobInfo(job);
+    }
+
+    async cancelTraining(jobId: string): Promise<TrainingJob> {
+        const job = await this.service.fineTuning.jobs.cancel(jobId);
+        return jobInfo(job);
+    }
+
+    async getTrainingJob(jobId: string): Promise<TrainingJob> {
+        const job = await this.service.fineTuning.jobs.retrieve(jobId);
+        return jobInfo(job);
+    }
+
+    // ========= management API =============
+
+    async validateConnection(): Promise<boolean> {
+        try {
+            await this.service.models.list();
+            return true;
+        } catch (error) {
+            return false;
+        }
+    }
+
+    listTrainableModels(): Promise<AIModel<string>[]> {
+        return this._listModels((m) => supportFineTunning.has(m.id));
+    }
+
+    async listModels(): Promise<AIModel[]> {
+        return this._listModels();
+    }
+
+    async _listModels(filter?: (m: OpenAI.Models.Model) => boolean) {
+        let result = await this.service.models.list();
+        const models = filter ? result.data.filter(filter) : result.data;
+        return models.map((m) => ({
+            id: m.id,
+            name: m.id,
+            provider: this.provider,
+            owner: m.owned_by,
+            type: m.object === "model" ? ModelType.Text : ModelType.Unknown,
+        }));
+    }
+
+
+    async generateEmbeddings(content: string, model: string = "text-embedding-ada-002"): Promise<{ embeddings: number[], model: string; }> {
+        const res = await this.service.embeddings.create({
+            input: content,
+            model: model,
+        });
+
+        const embeddings = res.data[0].embedding;
+
+        if (!embeddings || embeddings.length === 0) {
+            throw new Error("No embedding found");
+        }
+
+        return { embeddings, model };
+    }
+
+}
+
+
+function jobInfo(job: OpenAI.FineTuning.Jobs.FineTuningJob): TrainingJob {
+    //validating_files`, `queued`, `running`, `succeeded`, `failed`, or `cancelled`.
+    const jobStatus = job.status;
+    let status = TrainingJobStatus.running;
+    let details: string | undefined;
+    if (jobStatus === 'succeeded') {
+        status = TrainingJobStatus.succeeded;
+    } else if (jobStatus === 'failed') {
+        status = TrainingJobStatus.failed;
+        details = job.error ? `${job.error.code} - ${job.error.message} ${job.error.param ? " [" + job.error.param + "]" : ""}` : "error";
+    } else if (jobStatus === 'cancelled') {
+        status = TrainingJobStatus.cancelled;
+    } else {
+        status = TrainingJobStatus.running;
+        details = jobStatus;
+    }
+    return {
+        id: job.id,
+        model: job.fine_tuned_model || undefined,
+        status,
+        details
+    }
+}