[GPT-J] How to get token embeddings ?

[ocordeiro]

Hi, I'm trying get vectors of token embeddings for contextual search proposes.
How to get this?

In python I do this:

import torch
from transformers import GPT2Tokenizer, GPT2Model
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2Model.from_pretrained('gpt2')

texto = "O Brasil é o país mais populoso da América Latina."

tokens = tokenizer.encode(texto, return_tensors='pt')

with torch.no_grad():
    embeddings = model(tokens)[0]

print(embeddings)

But, I'm not sure if my editing of main.cpp from GPT-J is correct

struct ggml_tensor * inpL = ggml_get_rows(ctx0, model.wte, embd);

for (int i = 0; i < n_embd; ++i) { 
     printf("%f, ", ((float *)inpL->data)[i]);
}

Here the complete code: https://gist.github.com/ocordeiro/eba67c2a1e6e2bf4fca702abb19e1ee7

[ocordeiro]

I achieved :D

token =  my
index = [616]
wte = [-0.004190, -0.008380, 0.020950, -0.025140,  ... 0.016206, -0.004051, -0.020257, 0.024309]


token =  my
index = [616]
wte = [0.000000, -0.018284, 0.018284, -0.007314,  ... 0.000000, 0.002982, -0.017891, -0.005964]


token =  dog
index = [3290]
wte = [0.000000, -0.018284, 0.018284, -0.007314,  ... 0.000000, 0.002982, -0.017891, -0.005964]

Report:

main: mem per token =    16484 bytes
main:     load time =  2174.95 ms
main:   sample time =     0.00 ms
main:  predict time =     0.34 ms / 0.08 ms per token
main:    total time =  2177.07 ms

[ggerganov]

Great! If you want, you can share the code here for future references

[ocordeiro]

Its here:
https://gist.githubusercontent.com/ocordeiro/eba67c2a1e6e2bf4fca702abb19e1ee7/raw/9497f5c979042d0e47f0e067df7e12555398d461/main.cpp

I think it's possible to reduce the loading time by loading only the model's embeddings tensor into memory.
I will try to do this.

[Ayushk4]

I believe that it is incorrect as you are printing the values before you run the computation graph:

    // run the computation
    ggml_build_forward_expand(&gf, inpL);
    ggml_graph_compute       (ctx0, &gf);

I believe this is when the graph is being executed, and you print before it.

[ocordeiro]

@Ayushk4 You are right. thanks for alerting.

[Ayushk4]

Sure.