WaledRashed24/comics

1"use server"2 3import { HfInference, HfInferenceEndpoint } from "@huggingface/inference"4import { LLMEngine, LLMPredictionFunctionParams } from "@/types"5import { createZephyrPrompt } from "@/lib/createZephyrPrompt"6 7export async function predict({8  systemPrompt,9  userPrompt,10  nbMaxNewTokens,11  // llmVendorConfig // <-- arbitrary/custom LLM models hosted on HF is not supported yet using the UI12}: LLMPredictionFunctionParams): Promise<string> {13 14  const hf = new HfInference(process.env.AUTH_HF_API_TOKEN)15 16  const llmEngine = `${process.env.LLM_ENGINE || ""}` as LLMEngine17  const inferenceEndpoint = `${process.env.LLM_HF_INFERENCE_ENDPOINT_URL || ""}`18  const inferenceModel = `${process.env.LLM_HF_INFERENCE_API_MODEL || ""}`19 20  let hfie: HfInferenceEndpoint = hf21 22  switch (llmEngine) {23    case "INFERENCE_ENDPOINT":24      if (inferenceEndpoint) {25       //  console.log("Using a custom HF Inference Endpoint")26        hfie = hf.endpoint(inferenceEndpoint)27      } else {28        const error = "No Inference Endpoint URL defined"29        console.error(error)30        throw new Error(error)31      }32      break;33    34    case "INFERENCE_API":35      if (inferenceModel) {36        // console.log("Using an HF Inference API Model")37      } else {38        const error = "No Inference API model defined"39        console.error(error)40        throw new Error(error)41      }42      break;43 44    default:45      const error = "Please check your Hugging Face Inference API or Inference Endpoint settings"46      console.error(error)47      throw new Error(error)48  }49    50  const api = llmEngine === "INFERENCE_ENDPOINT" ? hfie : hf51 52  let instructions = ""53  try {54    for await (const output of api.textGenerationStream({55      model: llmEngine === "INFERENCE_ENDPOINT" ? undefined : (inferenceModel || undefined),56      57      inputs: createZephyrPrompt([58        { role: "system", content: systemPrompt },59        { role: "user", content: userPrompt }60      ]) + "\n[{", // <-- important: we force its hand61 62      parameters: {63        do_sample: true,64        max_new_tokens: nbMaxNewTokens,65        return_full_text: false,66      }67    })) {68      instructions += output.token.text69      // process.stdout.write(output.token.text)70      if (71        instructions.includes("</s>") || 72        instructions.includes("<s>") ||73        instructions.includes("/s>") ||74        instructions.includes("[INST]") ||75        instructions.includes("[/INST]") ||76        instructions.includes("<SYS>") ||77        instructions.includes("<<SYS>>") ||78        instructions.includes("</SYS>") ||79        instructions.includes("<</SYS>>") ||80        instructions.includes("<|user|>") ||81        instructions.includes("<|end|>") ||82        instructions.includes("<|system|>") ||83        instructions.includes("<|assistant|>")84      ) {85        break86      }87    }88  } catch (err) {89    // console.error(`error during generation: ${err}`)90 91    // a common issue with Llama-2 might be that the model receives too many requests92    if (`${err}` === "Error: Model is overloaded") {93      instructions = ``94    }95  }96 97  // need to do some cleanup of the garbage the LLM might have gave us98  return (99    instructions100    .replaceAll("<|end|>", "")101    .replaceAll("<s>", "")102    .replaceAll("</s>", "")103    .replaceAll("/s>", "")104    .replaceAll("[INST]", "")105    .replaceAll("[/INST]", "") 106    .replaceAll("<SYS>", "")107    .replaceAll("<<SYS>>", "")108    .replaceAll("</SYS>", "")109    .replaceAll("<</SYS>>", "")110    .replaceAll("<|system|>", "")111    .replaceAll("<|user|>", "")112    .replaceAll("<|all|>", "")113    .replaceAll("<|assistant|>", "")114    .replaceAll('""', '"')115  )116}117