diff --git a/Cargo.toml b/Cargo.toml index 7665256..a6e00cb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ edition = "2024" [dependencies] anyhow = "1.0.102" +directories = "6.0.0" lazy_static = "1.5.0" ndarray = "0.17.2" ort = "2.0.0-rc.11" @@ -12,3 +13,4 @@ serde = "1.0.228" serde_json = "1.0.149" tempfile = "3.26.0" tokenizers = "0.22.2" +toml = "1.0.3" diff --git a/README.md b/README.md index 08c0292..c573c42 100644 --- a/README.md +++ b/README.md @@ -1,10 +1,9 @@ # SUIME -**SUIME** 是一个基于深度神经网络的高性能输入法引擎项目。它采用现代化的 **C/S (Client/Server)** 架构,将繁重的 AI 推理任务与轻量级的输入交互界面解耦。 +**SUIME** 是一个用于学习的基于深度神经网络的输入法引擎项目。 - 🚀 **核心智能**:基于 Rust 构建的高性能推理服务端,支持 INT8 量化 ONNX 模型。 - 🐧 **原生 Linux 支持**:提供基于 `fcitx5` 的高性能插件。 -- 🪟 **跨平台愿景**:架构设计原生支持 Windows (计划基于 Rime/Weasel) 和 macOS。 - 🔌 **语言无关协议**:基于 MessagePack 的通用通信协议,允许任意语言编写前端。 --- @@ -48,41 +47,39 @@ graph TD - 基于 `fcitx5` 框架开发的 C++ 插件。 - 直接通过 UDS 与服务端通信。 - 负责捕获按键、获取上下文、渲染候选词面板。 -- **Windows (中长期计划)**: - - **目标框架**:基于 **Rime (中州韵)** 的 Windows 前端 **Weasel (小狼毫)** 进行定制开发。 - - **实现方式**:修改 Weasel 源码,将其内部解码逻辑替换为通过网络调用 Rust 服务端。 - - *注:由于 Windows TSF 架构复杂性,直接开发原生 IME 难度极大,复用 Rime 生态是最佳路径。* -- **macOS (未来计划)**: - - 基于 **Rime** 的 macOS 前端 **Squirrel (鼠须管)** 进行类似定制。 - --- ## 📂 项目目录结构 +### 目录 + ```text SUIME/ ├── Cargo.toml # Rust 服务端配置 ├── src/ # Rust 服务端源码 +│ ├── config.rs # 配置文件解析与加载,配置文件放在user_config_dir(类似~/.config/suime/config.toml),无配置可直接生成默认配置 +│ ├── filter.rs # 对候选词进行转换和筛选,将OnnxModel::predict预测的结果进行转换,并根据用户输入的Request.pinyin进行匹配和筛选,返回候选字列表 +│ ├── vocabs.rs # id和汉字转化 +│ ├── personal.rs # 个性化词汇的登录、管理和查询(未来扩展) +│ ├── protocol.rs # 跨平台通用的请求/响应结构体 (MessagePack) │ ├── main.rs # 入口:根据 OS 自动选择监听模式 (UDS/TCP) │ ├── model.rs # ONNX 模型加载与推理封装 -│ ├── config.rs # 配置文件解析与加载,配置文件放在user_config_dir(类似~/.config/suime/config.toml),无配置可直接生成默认配置 -│ ├── tokenizer.rs # 分词器封装 -│ ├── protocol.rs # 跨平台通用的请求/响应结构体 (MessagePack) -│ └── socket.rs # 通信层抽象 (支持 UDS/TCP 条件编译) -│ └── personal.rs # 个性化词汇的登录管理和查询(未来扩展) +│ ├── socket.rs # 通信层抽象 +│ └── tokenizer.rs # 分词器封装 ├── fcitx5-ext/ # [Linux] fcitx5 插件源码 │ ├── CMakeLists.txt │ ├── src/ -│ │ ├── myengine.cpp # 输入法逻辑 +│ │ ├── suime.cpp # 输入法逻辑 │ │ ├── socket_client.cpp # Linux UDS 客户端实现 │ │ └── protocol.hpp # 与服务端一致的协议定义 -├── clients/ # [未来] 其他平台客户端 -│ ├── weasel-mod/ # (计划) Windows Rime 定制版 -│ └── squirrel-mod/ # (计划) macOS Rime 定制版 ├── assets/ # ONNX 模型文件 & 词表 └── README.md ``` +### 文件说明 + +- **filter.rs** 将OnnxModel::predict预测的结果转化为(汉字、拼音、权重),排除和用户输入的拼音不相符的汉字,比如用户输入的是shanghai,预测的结果为(上,商,尚,特,……),特将会被排除,为了简化,直接粗暴的将预测汉字的拼音首字母和用户输入的拼音首字母进行对比,相同视为相符,不相同视为不符。 + --- ## 🛠️ 技术细节与接口定义 @@ -134,28 +131,11 @@ pub struct Response { --- ## 🗺️ 开发与部署路线图 - -### 阶段一:Linux 原生完善 (Current) - [ ] 完成 Rust 服务端 UDS 通信与 ONNX 推理。 - [ ] 完成 fcitx5 插件开发与联调。 - [ ] 性能优化:引入线程池,减少上下文切换开销。 - [ ] 打包:提供 `.deb` / `.rpm` 安装包及 systemd 用户服务配置。 -### 阶段二:服务端跨平台重构 (Short Term) -- [ ] **通信层抽象**:重构 `socket.rs`,实现 UDS/TCP 自动切换逻辑。 -- [ ] **Windows 服务端验证**:在 Windows 上编译并运行 Rust 服务端,通过 TCP 暴露接口。 -- [ ] **macOS 服务端验证**:验证 macOS 下 UDS 兼容性及签名问题。 - -### 阶段三:Windows 客户端适配 (Mid-Long Term) -- [ ] **Rime/Weasel 调研**:深入分析 Weasel 源码,确定 Hook 点。 -- [ ] **原型开发**:Fork Weasel 项目,移除内置拼音引擎,植入 TCP 客户端模块。 -- [ ] **UI 适配**:确保候选词窗口在高分屏下的渲染效果。 -- [ ] **安装包制作**:制作 Windows 安装程序,包含 Rust 服务端 (后台服务) 和定制版 Weasel。 - -### 阶段四:macOS 客户端适配 (Future) -- [ ] 基于 Squirrel 进行类似 Weasel 的定制开发。 -- [ ] 适配 macOS 输入法沙盒机制 (可能需要将服务端作为 LaunchAgent 运行)。 - --- ## 🚀 快速开始 (Linux) diff --git a/src/config.rs b/src/config.rs index e69de29..92531d3 100644 --- a/src/config.rs +++ b/src/config.rs @@ -0,0 +1,77 @@ +// config.rs + +use serde::{Deserialize, Serialize}; +use std::fs; +use std::path::PathBuf; +use directories::ProjectDirs; + +#[derive(Debug, Serialize, Deserialize)] +pub struct Config { + /// 模型文件路径(绝对路径) + pub model_path: String, + /// 词表文件路径(绝对路径) + pub tokenizer_path: String, + /// Socket配置:对于Unix是路径,对于Windows是地址 + pub socket: SocketConfig, + /// 额外配置,为未来扩展预留 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub extra: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub enum SocketConfig { + UnixDomainSocket(String), + TcpSocket(String), +} + +impl Config { + /// 从配置文件加载配置,如果文件不存在,则生成并写入默认配置 + pub fn load() -> Result> { + if let Some(proj_dirs) = ProjectDirs::from("", "", "suime") { + let config_dir = proj_dirs.config_dir(); + if !config_dir.exists() { + fs::create_dir_all(config_dir)?; + } + let config_file = config_dir.join("config.toml"); + + if config_file.exists() { + let config_str = fs::read_to_string(&config_file)?; + let config: Config = toml::from_str(&config_str)?; + Ok(config) + } else { + let default_config = Config::default(); + let config_str = toml::to_string_pretty(&default_config)?; + fs::write(config_file, config_str)?; + Ok(default_config) + } + } else { + Err("无法确定配置目录".into()) + } + } +} + +impl Default for Config { + fn default() -> Self { + // 获取当前工作目录作为基准,计算绝对路径 + let base_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + let model_path = base_dir.join("assets/2026030201suinput_int8(acc80).onnx") + .to_string_lossy() + .to_string(); + let tokenizer_path = base_dir.join("assets/tokenizer.json") + .to_string_lossy() + .to_string(); + + #[cfg(unix)] + let socket = SocketConfig::UnixDomainSocket("/tmp/su-ime.sock".to_string()); + + #[cfg(windows)] + let socket = SocketConfig::TcpSocket("127.0.0.1:23333".to_string()); + + Config { + model_path, + tokenizer_path, + socket, + extra: None, // 默认无额外配置 + } + } +} diff --git a/src/main.rs b/src/main.rs index 2c030d4..3a154bf 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,22 +4,17 @@ use std::time::Instant; mod model; mod tokenizers; mod vocabs; +mod config; use tokenizers::HFTokenizer; // use crate::vocabs::Dictionary; use crate::model::OnnxModel; +use crate::config::Config; fn main() { - let mut tokenizer_json_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")); - tokenizer_json_path.push("assets"); - tokenizer_json_path.push("tokenizer.json"); - - let mut model_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")); - model_path.push("assets"); - model_path.push("20260228suinput_fp32.onnx"); - - let mut session = OnnxModel::new(model_path, 4).unwrap(); - let tokenizer = HFTokenizer::new(tokenizer_json_path).unwrap(); + let config = Config::load().unwrap() ; + let mut session = OnnxModel::new(config.model_path, 4).unwrap(); + let tokenizer = HFTokenizer::new(config.tokenizer_path).unwrap(); let start = Instant::now(); // 开始计时 let sample = tokenizer.gen_predict_sample("从北京到上", "hai").unwrap(); let _ = session.predict(sample).unwrap();