feat: add flags to enable GPU acceleration

Author: jkawamoto

src/bin/vsops.rs

@@ -11,6 +11,8 @@ use std::path::PathBuf;
 
 use anyhow::{anyhow, Result};
 use clap::{crate_name, Parser};
+use ct2rs::config::Config;
+use ct2rs::config::Device;
 use tokio::signal;
 use tokio::sync::oneshot;
 
@@ -45,6 +47,12 @@ struct Args {
     /// Specifies the path to the socket file.
     #[arg(long)]
     socket_file: Option<PathBuf>,
+    /// Enable GPU acceleration,
+    #[arg(long)]
+    cuda: bool,
+    /// Specifies the CUDA device ID to be used. Effective only when the `--cuda` flag is enabled.
+    #[arg(long, default_value = "0")]
+    cuda_device_id: i32,
 }
 
 #[tokio::main]
@@ -67,7 +75,14 @@ async fn main() -> Result<()> {
         Some(path) => SocketFile::with_path(path)?,
         None => SocketFile::new(crate_name!())?,
     };
-    let server = Server::new(model_dir)?;
+    let server = Server::new(
+        model_dir,
+        Config {
+            device: if args.cuda { Device::CUDA } else { Device::CPU },
+            device_indices: vec![args.cuda_device_id],
+            ..Default::default()
+        },
+    )?;
     server
         .serve(socket_file, async move {
             if let Err(e) = rx.await {

src/lib.rs

@@ -11,6 +11,7 @@ use std::path::{Path, PathBuf};
 
 use anyhow::{Error, Result};
 use ct2rs::auto::Tokenizer as AutoTokenizer;
+use ct2rs::config::Config;
 use ct2rs::{TranslationOptions, Translator};
 use tokio::net::{UnixListener, UnixStream};
 use tokio_stream::wrappers::UnixListenerStream;
@@ -60,9 +61,9 @@ pub struct Server {
 }
 
 impl Server {
-    pub fn new<P: AsRef<Path>>(model_path: P) -> Result<Self> {
+    pub fn new<P: AsRef<Path>>(model_path: P, config: Config) -> Result<Self> {
         Ok(Self {
-            inner: Translator::new(model_path, &Default::default())?,
+            inner: Translator::new(model_path, &config)?,
             options: TranslationOptions {
                 beam_size: 12,
                 use_vmap: true,

src/socket.rs

@@ -46,6 +46,9 @@ impl AsRef<Path> for SocketFile {
 
 impl Drop for SocketFile {
     fn drop(&mut self) {
+        if !self.path.exists() {
+            return;
+        }
         if let Err(e) = remove_file(&self.path) {
             println!("failed to remove socket file {}: {}", &self, e);
         }