diff --git a/gitnexus/src/core/embeddings/embedder.ts b/gitnexus/src/core/embeddings/embedder.ts index cac09743b5..f4d969348b 100644 --- a/gitnexus/src/core/embeddings/embedder.ts +++ b/gitnexus/src/core/embeddings/embedder.ts @@ -14,7 +14,12 @@ if (!process.env.ORT_LOG_LEVEL) { process.env.ORT_LOG_LEVEL = '3'; } -import { pipeline, env, type FeatureExtractionPipeline } from '@huggingface/transformers'; +import { + pipeline, + env, + type FeatureExtractionPipeline, + type ProgressInfo, +} from '@huggingface/transformers'; import { existsSync } from 'fs'; import { execFileSync } from 'child_process'; import { join, dirname } from 'path'; @@ -171,13 +176,13 @@ export const initEmbedder = async ( } const progressCallback = onProgress - ? (data: any) => { + ? (data: ProgressInfo) => { const progress: ModelProgress = { - status: data.status || 'progress', - file: data.file, - progress: data.progress, - loaded: data.loaded, - total: data.total, + status: data.status === 'progress_total' ? 'progress' : data.status, + file: 'file' in data ? data.file : undefined, + progress: 'progress' in data ? data.progress : undefined, + loaded: 'loaded' in data ? data.loaded : undefined, + total: 'total' in data ? data.total : undefined, }; onProgress(progress); } @@ -204,7 +209,7 @@ export const initEmbedder = async ( embedderInstance = await withHfDownloadRetry( () => - (pipeline as any)('feature-extraction', finalConfig.modelId, { + pipeline('feature-extraction', finalConfig.modelId, { device: device, dtype: 'fp32', progress_callback: progressCallback, diff --git a/gitnexus/src/mcp/core/embedder.ts b/gitnexus/src/mcp/core/embedder.ts index ddd86ce5e2..f529f28054 100644 --- a/gitnexus/src/mcp/core/embedder.ts +++ b/gitnexus/src/mcp/core/embedder.ts @@ -74,7 +74,7 @@ export const initEmbedder = async (): Promise => { process.stderr.write = (() => true) as any; try { embedderInstance = await withHfDownloadRetry(() => - (pipeline as any)('feature-extraction', MODEL_ID, { + pipeline('feature-extraction', MODEL_ID, { device: device, dtype: 'fp32', session_options: { diff --git a/gitnexus/test/unit/hf-env.test.ts b/gitnexus/test/unit/hf-env.test.ts index 8905b11c9b..fa8f23bd34 100644 --- a/gitnexus/test/unit/hf-env.test.ts +++ b/gitnexus/test/unit/hf-env.test.ts @@ -10,7 +10,6 @@ import { withDownloadTimeout, withHfDownloadRetry, CIRCUIT_OPEN_TAG, - HF_DOWNLOAD_TIMEOUT_MS, HF_MAX_ATTEMPTS, HF_MAX_TIMEOUT_MS, HF_MAX_ATTEMPTS_CAP,