Fix embeddings extraction for all models

[skirodev]

Apply the fixed code for embeddings extraction to all models to avoid assertion errors. #288

[LLukas22]

Correct me if im wrong, until now we passed the output of the embedding layer as the model embedding?

[philpax]

Hmm yeah, I think this is correct (the embeddings should come from just before the LM head is evaluated to coalesce the LM output into logits). I'm on holiday so I can't check, but I assume embd is the input embeddings, not the output.

Happy to merge if you can validate my understanding.

[LLukas22]

In my oppinion this is also correct, but i'll test it when i get home. If it works i'll merge it.

[skirodev]

@LLukas22 Yes, exactly, we obtain the vector corresponding to the last token as the embeddings for the whole sentence. The length of this embeddings is a fixed value equal to n_embd. For the LLaMA 7B model, its length is 4096, while for OpenAI's text-embedding-ada-002 model, its length is 1536.

[LLukas22]

Good catch, when i get home from Work i'll test it against the rustformers HF models, but i'm 99% sure it will work.

[Jeadie]

Confirmed that this is a fix for #288

[LLukas22]

This works, but for gpt-j and gpt-neox based models im getting a process didn't exit successfully: target\release\examples\embeddings.exe gptj C:\Users\Lu.Kreuss\Downloads\gpt-j-6b-q4_0-ggjt.bin -r EleutherAI/gpt-j-6b (exit code: 0xc0000005, STATUS_ACCESS_VIOLATION) error. Tested with rustformers/gpt-j-ggml and rustformers/redpajama-ggml.

Im also getting very poor embeddings this way, compared to some BERT models i have lying around. Maybe we should perform some sort of pooling on the embeddings of all tokens. The SGPT paper uses weighted mean pooling, where more recent tokens have a stronger impact on the produced embedding than older tokens. Maybe this would improve the quality of the embeddings?

[skirodev]

This works, but for gpt-j and gpt-neox based models im getting a process didn't exit successfully: target\release\examples\embeddings.exe gptj C:\Users\Lu.Kreuss\Downloads\gpt-j-6b-q4_0-ggjt.bin -r EleutherAI/gpt-j-6b (exit code: 0xc0000005, STATUS_ACCESS_VIOLATION) error. Tested with rustformers/gpt-j-ggml and rustformers/redpajama-ggml.

This is because the value of "context_size" set in the model parameters was too large, so I rewrote the embeddings example.

Im also getting very poor embeddings this way, compared to some BERT models i have lying around. Maybe we should perform some sort of pooling on the embeddings of all tokens. The SGPT paper uses weighted mean pooling, where more recent tokens have a stronger impact on the produced embedding than older tokens. Maybe this would improve the quality of the embeddings?

This might be a viable approach, but I am not sure if it is necessary to use a model with a large number of parameters to generate embeddings?

[LLukas22]

This is because the value of "context_size" set in the model parameters was too large, so I rewrote the embeddings example.

o_O I feel kinda stupid for missing this, good job now everything works as expected 👍

This might be a viable approach, but I am not sure if it is necessary to use a model with a large number of parameters to generate embeddings?

It makes sense if you don't want to load a second smaller model to perform the embedding task. But i'm gooing to create another issue for this. It's not part of this PR.