feat: Added encode_single function (#18)

Author: Pringled

README.md

@@ -107,7 +107,7 @@ The compiled binary installed via `cargo install` is significantly faster (often
 
 *   **Encode a single sentence:**
     ```shell
-    model2vec-rs encode "Hello world" "minishlab/potion-base-8M"
+    model2vec-rs encode-single "Hello world" "minishlab/potion-base-8M"
     ```
     Embeddings will be printed to the console in JSON format. This command should take less than 0.1s to execute.
 

src/main.rs

@@ -26,11 +26,21 @@ enum Commands {
         #[arg(short, long)]
         output: Option<String>,
     },
+    /// Encode a single sentence  
+    EncodeSingle {
+        /// The sentence to embed
+        sentence: String,
+        /// HF repo ID or local dir
+        model: String,
+        #[arg(short, long)]
+        output: Option<String>,
+    },
 }
 
 fn main() -> Result<()> {
     let cli = Cli::parse();
     match cli.cmd {
+        // Encode multiple sentences from a file or input string
         Commands::Encode { input, model, output } => {
             let texts = if Path::new(&input).exists() {
                 std::fs::read_to_string(&input)?
@@ -52,6 +62,19 @@ fn main() -> Result<()> {
                 println!("{:?}", embs);
             }
         }
+        // Encode a single sentence
+        Commands::EncodeSingle { sentence, model, output } => {
+            let m = StaticModel::from_pretrained(&model, None, None, None)?;
+            let embedding = m.encode_single(&sentence);
+
+            if let Some(path) = output {
+                let file = File::create(path).context("creating output file failed")?;
+                serde_json::to_writer(BufWriter::new(file), &embedding)
+                    .context("writing JSON failed")?;
+            } else {
+                println!("{embedding:#?}");
+            }
+        }
     }
     Ok(())
 }

src/model.rs

@@ -197,6 +197,11 @@ impl StaticModel {
         self.encode_with_args(sentences, Some(512), 1024)
     }
 
+    // / Encode a single sentence into a vector
+    pub fn encode_single(&self, sentence: &str) -> Vec<f32> {
+        self.encode(&[sentence.to_string()]).into_iter().next().unwrap_or_default()
+    }
+
     /// Mean-pool a single token-ID list into a vector
     fn pool_ids(&self, ids: Vec<u32>) -> Vec<f32> {
         let mut sum = vec![0.0; self.embeddings.ncols()];

tests/test_model.rs

@@ -72,6 +72,35 @@ fn test_encode_empty_sentence() {
     assert!(vec.iter().all(|&x| x == 0.0), "All entries should be zero");
 }
 
+/// Test that encoding a single sentence returns the correct shape
+#[test]
+fn test_encode_single() {
+    let model = load_test_model();
+    let sentence = "hello world";
+
+    // Single-sentence helper → 1-D
+    let one_d = model.encode_single(sentence);
+
+    // Batch call with a 1-element slice → 2-D wrapper
+    let two_d = model.encode(&[sentence.to_string()]);
+
+    // Shape assertions
+    assert!(
+        !one_d.is_empty(),
+        "encode_single must return a non-empty 1-D vector"
+    );
+    assert_eq!(
+        two_d.len(),
+        1,
+        "encode(&[..]) should wrap the result in a Vec with length 1"
+    );
+    assert_eq!(
+        two_d[0].len(),
+        one_d.len(),
+        "inner vector dimensionality should match encode_single output"
+    );
+}
+
 /// Test override of `normalize` flag in from_pretrained
 #[test]
 fn test_normalization_flag_override() {