rust-bert/examples/sentence_embeddings.rs
Romain Leroux 4d8a298586
Add sbert implementation for inference (#250)
* Add sbert implementation for inference

* Fix clippy warnings

* Refactor sentence embeddings into a dedicated pipeline

* Add output_attentions and output_hidden_states to T5Config

* Add sbert implementation for inference

* Fix clippy warnings

* Refactor sentence embeddings into a dedicated pipeline

* Add output_attentions and output_hidden_states to T5Config

* Improve sentence_embeddings implementation

* Dedicated tokenizer config for strip_accents and add_prefix_space

* Rename forward to encode_as_tensor

* Remove _conf from Dense layer

* Add sentence embeddings docs

* Addition of remote resources and tests update

* Merge feature branch and fix doctests

* Add SentenceEmbeddingsBuilder<Remote> and improve remote resources

* Use tch::no_grad in sentence embeddings

* Updated changelog, registration of sentence embeddings integration tests

Co-authored-by: Guillaume Becquin <guillaume.becquin@gmail.com>
2022-06-21 20:24:09 +01:00

18 lines
535 B
Rust

use rust_bert::pipelines::sentence_embeddings::{
SentenceEmbeddingsBuilder, SentenceEmbeddingsModelType,
};
fn main() -> anyhow::Result<()> {
// Set-up sentence embeddings model
let model = SentenceEmbeddingsBuilder::remote(SentenceEmbeddingsModelType::AllMiniLmL12V2)
.create_model()?;
// Define input
let sentences = ["this is an example sentence", "each sentence is converted"];
// Generate Embeddings
let embeddings = model.encode(&sentences)?;
println!("{:?}", embeddings);
Ok(())
}