Download command for models

Models are still not being published, so this fails downloading. Will
work in the future hopufully.

Cargo.toml

@@ -22,6 +22,9 @@ strum = { version = "0.25", features = ["derive"] }
 strum_macros = "0.25"
 wyhash2 = "0.2.1"
 pyo3 = { version = "0.22", features = ["gil-refs"] }
+reqwest = { version = "0.12", features = ["stream"] }
+tokio = { version = "1", features = ["io-util", "rt-multi-thread", "signal"] }
+futures-util = "0.3"
 
 [dev-dependencies]
 test-log = "0.2.15"

pyproject.toml

@@ -15,6 +15,7 @@ dynamic = ["version"]
 [project.scripts]
 heli = "heli_otr:cli_run"
 heli-convert = "heli_otr:cli_convert"
+heli-download = "heli_otr:cli_download"
 
 [tool.maturin]
 features = ["pyo3/extension-module"]

src/lib.rs

@@ -15,10 +15,13 @@ use crate::identifier::Identifier;
 pub mod languagemodel;
 pub mod identifier;
 pub mod lang;
+mod utils;
 
+const WORDMODEL_FILE: &str = "wordmodel.bin";
+const CHARMODEL_FILE: &str = "charmodel.bin";
 
 // Call python interpreter and obtain python path of our module
-pub fn pythonpath() -> PyResult<String> {
+pub fn module_path() -> PyResult<String> {
     let mut path = String::new();
     Python::with_gil(|py| {
         // Instead of hardcoding the module name, obtain it from the crate name at compile time
@@ -33,13 +36,13 @@ pub fn pythonpath() -> PyResult<String> {
 }
 
 pub fn load_models(modelpath: &str) -> (Model, Model) {
-    let grampath = format!("{modelpath}/charmodel.bin");
+    let grampath = format!("{modelpath}/{CHARMODEL_FILE}");
     let char_handle = thread::spawn(move || {
         let path = Path::new(&grampath);
         Model::from_bin(path)
     });
 
-    let wordpath = format!("{modelpath}/wordmodel.bin");
+    let wordpath = format!("{modelpath}/{WORDMODEL_FILE}");
     let word_handle = thread::spawn(move || {
         let path = Path::new(&wordpath);
         Model::from_bin(path)
@@ -60,7 +63,7 @@ pub struct PyIdentifier {
 impl PyIdentifier {
     #[new]
     fn new() -> Self {
-        let modulepath = pythonpath().expect("Error loading python module path");
+        let modulepath = module_path().expect("Error loading python module path");
         let (charmodel, wordmodel) = load_models(&modulepath);
         let identifier = Identifier::new(
             Arc::new(charmodel),
@@ -86,7 +89,7 @@ impl PyIdentifier {
 #[pyfunction]
 pub fn cli_run() -> PyResult<()> {
     env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();
-    let modulepath = pythonpath().expect("Error loading python module path");
+    let modulepath = module_path().expect("Error loading python module path");
     let (charmodel, wordmodel) = load_models(&modulepath);
     let mut identifier = Identifier::new(
             Arc::new(charmodel),
@@ -101,10 +104,31 @@ pub fn cli_run() -> PyResult<()> {
     Ok(())
 }
 
+#[pyfunction]
+pub fn cli_download() -> PyResult<()> {
+    env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();
+    let modulepath = module_path().expect("Error loading python module path");
+    let url = format!(
+        "https://github.com/ZJaume/heli-otr/releases/download/v{}",
+        env!("CARGO_PKG_VERSION"));
+
+    utils::download_file(
+        &format!("{url}/{WORDMODEL_FILE}"),
+        &format!("{modulepath}/{WORDMODEL_FILE}")
+    ).unwrap();
+    utils::download_file(
+        &format!("{url}/{CHARMODEL_FILE}"),
+        &format!("{modulepath}/{CHARMODEL_FILE}")
+    ).unwrap();
+    info!("Finished");
+
+    Ok(())
+}
+
 #[pyfunction]
 pub fn cli_convert() -> PyResult<()> {
     env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();
-    let modulepath = pythonpath().expect("Error loading python module path");
+    let modulepath = module_path().expect("Error loading python module path");
     debug!("Module path found: {}", modulepath);
     let modelpath = Path::new("./LanguageModels");
 
@@ -127,6 +151,7 @@ pub fn cli_convert() -> PyResult<()> {
 fn heli_otr(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(cli_run))?;
     m.add_wrapped(wrap_pyfunction!(cli_convert))?;
+    m.add_wrapped(wrap_pyfunction!(cli_download))?;
     m.add_class::<PyIdentifier>()?;
     // m.add_class::<PyLang>()?;
 

src/utils.rs

@@ -0,0 +1,48 @@
+use std::process::exit;
+
+use log::{info, debug, error};
+use tokio::io::AsyncWriteExt;
+use tokio::runtime::Runtime;
+use futures_util::StreamExt;
+use reqwest;
+
+
+// Run a tokio task that listens for ctrl+c
+async fn run_cancel_handler() {
+    tokio::spawn(async move {
+        match tokio::signal::ctrl_c().await {
+            Ok(_) => {
+                info!("Received Ctrl+C, terminating immediately.");
+                exit(1);
+            }
+            Err(e) => error!("Error listening for SIGINT: {}", e),
+        };
+    });
+}
+
+// Download a file to a path
+pub fn download_file(url: &str, filepath: &str) -> Result<(), Box<dyn std::error::Error>> {
+    let runtime = Runtime::new()?;
+    runtime.block_on(async {
+        info!("Downloading file from '{url}'");
+        run_cancel_handler().await;
+        // Create a download stream
+        let response = reqwest::get(url).await?;
+        let status = response.status();
+        debug!("Response status: {}", status);
+        if !status.is_success() {
+            error!("Could not download file, HTTP status code: {status}");
+            exit(1);
+        }
+
+        let mut response_stream = response.bytes_stream();
+        let mut outfile = tokio::fs::File::create(filepath).await?;
+
+        debug!("Writing file to '{filepath}'");
+        // asyncronously write to the file every piece of bytes that come from the stream
+        while let Some(bytes) = response_stream.next().await {
+            outfile.write_all(&bytes?).await?;
+        }
+        Ok(())
+    })
+}