Add text embedder

flash_examples/text_embedder.py

@@ -27,9 +27,10 @@
 )
 
 # 2. Load a previously trained SentenceEmbedder
-model = SentenceEmbedder(backbone="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+model = SentenceEmbedder(backbone="sentence-transformers/all-MiniLM-L6-v2")
 
 # 3. Generate embeddings for the first 3 graphs
 trainer = flash.Trainer(gpus=torch.cuda.device_count())
 predictions = trainer.predict(model, datamodule=datamodule)
 print(predictions)
+

tests/text/embedding/test_model.py

@@ -32,7 +32,7 @@
 
 # ==============================
 
-TEST_BACKBONE = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"  # super small model for testing
+TEST_BACKBONE = "sentence-transformers/all-MiniLM-L6-v2"  # super small model for testing
 model = SentenceEmbedder(backbone=TEST_BACKBONE)