test emb

config.json

@@ -1,10 +1,11 @@
 {
   "architectures": [
-    "AriaForCausalLM"
+    "AriaForSequenceEmbeddings"
   ],
   "bos_token_id": 0,
   "eos_token_id": 1,
   "hidden_size": 1536,
+  "embedding_size": 512,
   "intermediate_size": 6144,
   "max_position_embeddings": 8192,
   "model_type": "aria",

modeling_aria.py

@@ -617,7 +617,7 @@ class AriaForCausalLM(AriaPreTrainedModel, GenerationMixin):
         )
 
 
-class AriaForSequenceEmbeddings(AriaPreTrainedModel):
+class AriaForSequenceEmbedding(AriaPreTrainedModel):
     """Transformer decoder embedding head for contrastive learning.
 
     Args:
@@ -744,5 +744,5 @@ __all__ = [
     "AriaModel",
     "TransformerBlock",
     "AriaForCausalLM",
-    "AriaForSequenceEmbeddings",
+    "AriaForSequenceEmbedding",
 ]