libffm 0.2.0

Field-aware factorization machines in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

use clap::{ColorChoice, Parser};
use libffm::Model;
use std::error::Error;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use std::process;

#[derive(Debug, Parser)]
#[command(name = "ffm-predict", version, color = ColorChoice::Never)]
struct Args {
    #[arg(value_parser)]
    test_file: PathBuf,

    #[arg(value_parser)]
    model_file: PathBuf,

    #[arg(value_parser)]
    output_file: PathBuf,
}

fn predict(test_path: &PathBuf, model_path: &PathBuf, output_path: &PathBuf) -> Result<(), Box<dyn Error>> {
    let model = Model::load(model_path)?;
    let (predictions, loss) = model.predict(test_path)?;

    let f_out = File::create(output_path)?;
    let mut writer = BufWriter::new(f_out);
    for prediction in predictions {
        writer.write(format!("{:.6}\n", prediction).as_bytes())?;
    }
    writer.flush()?;

    println!("logloss = {:.5}", loss);

    Ok(())
}

fn main() {
    let args = Args::parse();

    if let Err(err) = predict(&args.test_file, &args.model_file, &args.output_file) {
        println!("{}", err);
        process::exit(1);
    }
}