Failure to download Several Embeddings Models

[narner]

I'm trying to use nomic_text_v1 and nomic_text_v1_5, and getting the same error when trying to download both:

 UpdateError
  ▿ keyNotFound : 2 elements
    - base : "Linear"
    - key : "weight

Seems similar to https://github.com/ml-explore/mlx-swift-examples/issues/214

I'm on latest main when adding this package to my project

[davidkoski]

It is missing the weights for pooler in the NomicBertModel:

    public init(
        _ config: NomicBertConfiguration, pooler: Bool = true,
        lmHead: Bool = false
    ) {
...
        if pooler {
            self.pooler = Linear(config.embedDim, config.embedDim)
        } else {
            self.pooler = nil
        }

That doesn't come from configuration so I am not sure how you know to set it or not.

As a workaround you can change that parameter to false and it should load.

@anishbasu do you know how this is supposed to be configured?

[narner]

Do you mean forking it and manually changing that?

For reference, this is how I have I set up currently

private var modelContainer: ModelContainer?
private let modelConfig = ModelConfiguration.nomic_text_v1

let loadTask = Task {
     try await MLXEmbedders.loadModelContainer(hub: hubApi, configuration: modelConfig)
}
            
// Set a timeout for model loading
let container = try await withTimeout(seconds: 120) {
    try await loadTask.value
}

self.modelContainer = container

[davidkoski]

I was just thinking of changing it locally -- you can reference a local checkout in Xcode. If forking and making the change is easier, go for it.

I don't know if this is masking another issue.

[anishbasu]

Seems that they removed the pooler from the weights; this was a configurable attribute (and set to true) in the original code: https://huggingface.co/nomic-ai/nomic-embed-text-v1/blob/ea3f22e7569289f7442010b7e7f2f3e46b22e480/modeling_hf_nomic_bert.py but this was not exposed in the configuration as @davidkoski pointed out.

For now a hotfix can just be setting the pooler arg to false.

[narner]

Going to update this issue to reflect that I'm seeing errors while trying to download several other embeddings models, too:

• Nomic Embed Text v1 — Key weight not found in Linear
• Nomic Embed Text v1.5 — Key bias not found in Linear
• Snowflake Arctic EmbedXS — Key weight not found in Linear
• BGE M3 — Key weight not found in Linear
• Multilingual EG Small — The tokenizer type ‘XLMRobertaTokenizer’ is not supported
• Snowflake Arctic Embed Large — Key bias not found in Linear

[aPaleBlueDot]

@narner Are there at least some embedding models which worked?

[narner]

@narner Are there at least some embedding models which worked?

Yes, these ones all worked fine:

• BGE Micro v2
• GTE Tiny
• MiniLM L6 v2
• MiniLM L12 v2
• BGE Small EN v1.5
• BGE Base EN v1.5
• BGE Large EN v1.5
• Mixedbread Large v1

[aPaleBlueDot]

Thanks @narner , that's still a good set of choices. Do you use a reranker, if so with which other framework? It seems llama.cpp has the broadest model support, and it should play well with Swift, but I haven't tried it yet. [Sorry, if this is off topic for this thread.]

[narner]

@aPaleBlueDot yeah these are all working great. For context I'm building a Swift package for the Chroma vector database; and using this repo as a sub-dependency for leveraging some local embeddings models

[aPaleBlueDot]

@narner Nice! ping me upon launching it