20 lines
746 B
Python
20 lines
746 B
Python
import mlx.core as mx
|
|
from mlx import nn
|
|
|
|
from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
|
|
|
|
|
|
class CLIPEmbeddings(nn.Module):
|
|
def __init__(self, dims: int):
|
|
super().__init__()
|
|
self.position_embedding = nn.Embedding(num_embeddings=TokenizerCLIP.MAX_TOKEN_LENGTH, dims=dims)
|
|
self.token_embedding = nn.Embedding(num_embeddings=49408, dims=dims)
|
|
|
|
def __call__(self, tokens: mx.array) -> mx.array:
|
|
seq_length = tokens.shape[-1]
|
|
position_ids = mx.arange(0, seq_length).reshape(1, seq_length)
|
|
inputs_embeds = self.token_embedding(tokens)
|
|
position_embeddings = self.position_embedding(position_ids)
|
|
embeddings = inputs_embeds + position_embeddings
|
|
return embeddings
|