Update README.md

README.md

@@ -31,8 +31,21 @@ import mlx.core as mx
 
 model, tokenizer = load("mlx-community/embeddinggemma-300m-bf16")
 
-# For text embeddings
-output = generate(model, processor, texts=["I like grapes", "I like fruits"])
+
+# For text embedding
+sentences = [
+    "task: sentence similarity | query: Nothing really matters.",
+    "task: sentence similarity | query: The dog is barking.",
+    "task: sentence similarity | query: The dog is barking.",
+]
+
+encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='mlx')
+
+# Compute token embeddings
+input_ids = encoded_input['input_ids']
+attention_mask = encoded_input['attention_mask']
+output = model(input_ids, attention_mask)
+
 embeddings = output.text_embeds  # Normalized embeddings
 
 # Compute dot product between normalized embeddings
@@ -42,4 +55,21 @@ print("Similarity matrix between texts:")
 print(similarity_matrix)
 
 
+# You can use these task-specific prefixes for different tasks
+task_prefixes = {
+    "BitextMining": "task: search result | query: ",
+    "Clustering": "task: clustering | query: ",
+    "Classification": "task: classification | query: ",
+    "MultilabelClassification": "task: classification | query: ",
+    "PairClassification": "task: sentence similarity | query: ",
+    "InstructionRetrieval": "task: code retrieval | query: ",
+    "Reranking": "task: search result | query: ",
+    "Retrieval": "task: search result | query: ",
+    "Retrieval-query": "task: search result | query: ",
+    "Retrieval-document": "title: none | text: ",
+    "STS": "task: sentence similarity | query: ",
+    "Summarization": "task: summarization | query: ",
+    "document": "title: none | text: "
+}
+
 ```