Return numpy arrays

Author: danieldk

Cargo.toml

@@ -16,4 +16,5 @@ features = ["extension-module"]
 failure = "0.1"
 finalfusion = "0.5"
 libc = "0.2"
+numpy = "0.5"
 toml = "0.4"

src/lib.rs

@@ -9,6 +9,7 @@ use failure::Error;
 use finalfusion::metadata::Metadata;
 use finalfusion::prelude::*;
 use finalfusion::similarity::*;
+use numpy::{IntoPyArray, PyArray1};
 use pyo3::class::{basic::PyObjectProtocol, iter::PyIterProtocol};
 use pyo3::exceptions;
 use pyo3::prelude::*;
@@ -145,7 +146,7 @@ impl PyEmbeddings {
     ///
     /// If the word is not known, its representation is approximated
     /// using subword units.
-    fn embedding(&self, word: &str) -> PyResult<Vec<f32>> {
+    fn embedding(&self, word: &str) -> PyResult<Py<PyArray1<f32>>> {
         let embeddings = self.embeddings.borrow();
 
         use EmbeddingsWrap::*;
@@ -155,7 +156,10 @@ impl PyEmbeddings {
         };
 
         match embedding {
-            Some(embedding) => Ok(embedding.as_view().to_vec()),
+            Some(embedding) => {
+                let gil = pyo3::Python::acquire_gil();
+                Ok(embedding.into_owned().into_pyarray(gil.python()).to_owned())
+            }
             None => Err(exceptions::KeyError::py_err("Unknown word and n-grams")),
         }
     }
@@ -306,7 +310,7 @@ impl PyIterProtocol for PyEmbeddingIterator {
         Ok(slf.into())
     }
 
-    fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<(String, Vec<f32>)>> {
+    fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<(String, Py<PyArray1<f32>>)>> {
         let slf = &mut *slf;
 
         let embeddings = slf.embeddings.borrow();
@@ -327,7 +331,11 @@ impl PyIterProtocol for PyEmbeddingIterator {
 
             slf.idx += 1;
 
-            Ok(Some((word, embed.as_view().to_vec())))
+            let gil = pyo3::Python::acquire_gil();
+            Ok(Some((
+                word,
+                embed.into_owned().into_pyarray(gil.python()).to_owned(),
+            )))
         } else {
             Ok(None)
         }