21 lines
713 B
Python
21 lines
713 B
Python
import mlx.core as mx
|
|
from mlx import nn
|
|
|
|
from mflux.models.text_encoder.t5_encoder.t5_block import T5Block
|
|
from mflux.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
|
|
|
|
|
|
class T5Encoder(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.shared = nn.Embedding(num_embeddings=32128, dims=4096)
|
|
self.t5_blocks = [T5Block(i) for i in range(24)]
|
|
self.final_layer_norm = T5LayerNorm()
|
|
|
|
def forward(self, tokens: mx.array):
|
|
hidden_states = self.shared(tokens)
|
|
for block in self.t5_blocks:
|
|
hidden_states = block.forward(hidden_states)
|
|
hidden_states = self.final_layer_norm.forward(hidden_states)
|
|
return hidden_states
|