@threaded/ai 1.0.28 → 1.0.30

package/dist/index.cjs

@@ -232,6 +232,7 @@ var addUsage = (existing, promptTokens, completionTokens, totalTokens) => ({
 // src/embed.ts
 var modelCache = /* @__PURE__ */ new Map();
 var embed = async (model2, text, config) => {
+  const isBatch = Array.isArray(text);
   if (model2.startsWith("openai/")) {
     const modelName = model2.replace("openai/", "");
     const apiKey = getKey("openai") || process.env.OPENAI_API_KEY;
@@ -258,6 +259,9 @@ var embed = async (model2, text, config) => {
       throw new Error(`OpenAI API error: ${error}`);
     }
     const data = await response.json();
+    if (isBatch) {
+      return data.data.map((d) => d.embedding);
+    }
     return data.data[0].embedding;
   }
   try {
@@ -269,6 +273,14 @@ var embed = async (model2, text, config) => {
       modelCache.set(model2, extractor2);
     }
     const extractor = modelCache.get(model2);
+    if (isBatch) {
+      const results = [];
+      for (const t of text) {
+        const result2 = await extractor(t, { pooling: "mean", normalize: true });
+        results.push(Array.from(result2.data));
+      }
+      return results;
+    }
     const result = await extractor(text, { pooling: "mean", normalize: true });
     return Array.from(result.data);
   } catch (error) {
@@ -1127,7 +1139,7 @@ var callHuggingFace = async (config, ctx) => {
   const { pipeline } = await import("@huggingface/transformers");
   if (!modelCache2.has(model2)) {
     const generator2 = await pipeline("text-generation", model2, {
-      dtype: "q4f16"
+      dtype: "q4"
     });
     modelCache2.set(model2, generator2);
   }