phrasekit 0.2.0-x86_64-linux

data/ext/phrasekit/src/bin/phrasekit_tag.rs

@@ -0,0 +1,320 @@
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::env;
+use std::fs::File;
+use std::io::{BufRead, BufReader, BufWriter, Write};
+use std::path::Path;
+use std::process;
+
+#[path = "../payload.rs"]
+mod payload;
+
+use payload::Payload;
+
+#[derive(Debug, Deserialize)]
+struct TagConfig {
+    automaton_path: String,
+    payloads_path: String,
+    manifest_path: String,
+    vocab_path: String,
+    #[serde(default = "default_policy")]
+    policy: String,
+    #[serde(default = "default_max_spans")]
+    max_spans: usize,
+    #[serde(default = "default_label")]
+    label: String,
+}
+
+fn default_policy() -> String {
+    "leftmost_longest".to_string()
+}
+
+fn default_max_spans() -> usize {
+    100
+}
+
+fn default_label() -> String {
+    "PHRASE".to_string()
+}
+
+#[derive(Debug, Deserialize)]
+struct InputDocument {
+    doc_id: String,
+    tokens: Vec<String>,
+}
+
+#[derive(Debug, Serialize)]
+struct OutputDocument {
+    doc_id: String,
+    tokens: Vec<String>,
+    spans: Vec<Span>,
+}
+
+#[derive(Debug, Serialize)]
+struct Span {
+    start: usize,
+    end: usize,
+    phrase_id: u32,
+    label: String,
+}
+
+#[derive(Debug, Deserialize)]
+struct Vocabulary {
+    tokens: HashMap<String, u32>,
+    special_tokens: HashMap<String, u32>,
+}
+
+#[derive(Debug)]
+struct TaggingStats {
+    documents: usize,
+    total_spans: usize,
+    docs_with_spans: usize,
+}
+
+fn encode_tokens(tokens: &[String], vocab: &Vocabulary) -> Vec<u32> {
+    let unk_id = vocab.special_tokens.get("<UNK>").copied().unwrap_or(0);
+
+    tokens
+        .iter()
+        .map(|token| {
+            let normalized = token.to_lowercase();
+            vocab.tokens.get(&normalized).copied().unwrap_or(unk_id)
+        })
+        .collect()
+}
+
+fn tag_corpus(
+    corpus_path: &str,
+    config: &TagConfig,
+    output_path: &str,
+) -> Result<TaggingStats, Box<dyn std::error::Error>> {
+    println!("🏷️  PhraseKit Corpus Tagging");
+    println!("════════════════════════════════════════");
+    println!("Corpus:     {}", corpus_path);
+    println!("Config:     <config>");
+    println!("Output:     {}", output_path);
+    println!();
+
+    println!("📚 Loading matcher artifacts...");
+
+    let vocab_data = std::fs::read_to_string(&config.vocab_path)?;
+    let vocab: Vocabulary = serde_json::from_str(&vocab_data)?;
+    println!("  ✓ Loaded vocabulary ({} tokens)", vocab.tokens.len());
+
+    use daachorse::DoubleArrayAhoCorasick;
+    let automaton_bytes = std::fs::read(&config.automaton_path)?;
+    let (automaton, _): (DoubleArrayAhoCorasick<u32>, _) = unsafe {
+        DoubleArrayAhoCorasick::deserialize_unchecked(&automaton_bytes)
+    };
+    println!("  ✓ Loaded automaton");
+
+    let payloads_file = File::open(&config.payloads_path)?;
+    let payloads_reader = BufReader::new(payloads_file);
+    let payloads = payload::load_payloads(payloads_reader)?;
+    println!("  ✓ Loaded {} phrase payloads", payloads.len());
+
+    #[derive(Debug, Deserialize)]
+    struct Manifest {
+        separator_id: u32,
+    }
+
+    let manifest_data = std::fs::read_to_string(&config.manifest_path)?;
+    let manifest: Manifest = serde_json::from_str(&manifest_data)?;
+    println!("  ✓ Loaded manifest");
+    println!();
+
+    println!("🔍 Tagging documents...");
+
+    let corpus_file = File::open(corpus_path)?;
+    let corpus_reader = BufReader::new(corpus_file);
+
+    let output_file = File::create(output_path)?;
+    let mut output_writer = BufWriter::new(output_file);
+
+    let mut stats = TaggingStats {
+        documents: 0,
+        total_spans: 0,
+        docs_with_spans: 0,
+    };
+
+    for line in corpus_reader.lines() {
+        let line = line?;
+        if line.trim().is_empty() {
+            continue;
+        }
+
+        let doc: InputDocument = serde_json::from_str(&line)?;
+
+        let token_ids = encode_tokens(&doc.tokens, &vocab);
+
+        let separator = manifest.separator_id;
+        let mut bytes = Vec::with_capacity(token_ids.len() * 5);
+        for &token_id in &token_ids {
+            bytes.extend_from_slice(&token_id.to_le_bytes());
+            bytes.extend_from_slice(&separator.to_le_bytes());
+        }
+
+        #[derive(Debug, Clone, Copy)]
+        struct Match {
+            start: usize,
+            end: usize,
+            phrase_id: u32,
+        }
+
+        let mut matches: Vec<Match> = automaton
+            .find_overlapping_iter(&bytes)
+            .filter_map(|m| {
+                let pattern_id = m.value() as usize;
+                let start_token = m.start() / 8;
+                let end_token = (m.end() + 7) / 8;
+
+                payloads.get(pattern_id).map(|payload| Match {
+                    start: start_token,
+                    end: end_token,
+                    phrase_id: payload.phrase_id,
+                })
+            })
+            .collect();
+
+        if config.policy == "leftmost_longest" {
+            matches.sort_by_key(|m| (m.start, std::cmp::Reverse(m.end)));
+
+            let mut resolved = Vec::new();
+            let mut covered_end = 0;
+
+            for m in matches {
+                if m.start >= covered_end {
+                    resolved.push(m);
+                    covered_end = m.end;
+                }
+            }
+
+            matches = resolved;
+        } else if config.policy == "leftmost_first" {
+            matches.sort_by_key(|m| m.start);
+
+            let mut resolved = Vec::new();
+            let mut covered_end = 0;
+
+            for m in matches {
+                if m.start >= covered_end {
+                    resolved.push(m);
+                    covered_end = m.end;
+                }
+            }
+
+            matches = resolved;
+        }
+
+        if matches.len() > config.max_spans {
+            matches.truncate(config.max_spans);
+        }
+
+        let spans: Vec<Span> = matches
+            .into_iter()
+            .map(|m| Span {
+                start: m.start,
+                end: m.end,
+                phrase_id: m.phrase_id,
+                label: config.label.clone(),
+            })
+            .collect();
+
+        stats.total_spans += spans.len();
+        if !spans.is_empty() {
+            stats.docs_with_spans += 1;
+        }
+
+        let output_doc = OutputDocument {
+            doc_id: doc.doc_id,
+            tokens: doc.tokens,
+            spans,
+        };
+
+        serde_json::to_writer(&mut output_writer, &output_doc)?;
+        writeln!(&mut output_writer)?;
+
+        stats.documents += 1;
+
+        if stats.documents % 1000 == 0 {
+            print!("\r  Processed {} documents...", stats.documents);
+            std::io::stdout().flush()?;
+        }
+    }
+
+    if stats.documents % 1000 != 0 {
+        println!("\r  ✓ Processed {} documents", stats.documents);
+    } else {
+        println!();
+        println!("  ✓ Processed {} documents", stats.documents);
+    }
+
+    output_writer.flush()?;
+
+    println!();
+    println!("✅ Tagging complete!");
+    println!();
+    println!("📈 Statistics:");
+    println!("  Documents:              {}", stats.documents);
+    println!("  Total spans:            {}", stats.total_spans);
+    println!("  Documents with spans:   {}", stats.docs_with_spans);
+    println!(
+        "  Avg spans per document: {:.2}",
+        if stats.documents > 0 {
+            stats.total_spans as f64 / stats.documents as f64
+        } else {
+            0.0
+        }
+    );
+
+    Ok(stats)
+}
+
+fn main() {
+    let args: Vec<String> = env::args().collect();
+
+    if args.len() != 4 {
+        eprintln!("Usage: {} <corpus.jsonl> <config.json> <output.jsonl>", args[0]);
+        eprintln!();
+        eprintln!("Arguments:");
+        eprintln!("  corpus.jsonl  - Input corpus with pre-tokenized documents");
+        eprintln!("  config.json   - Tagging configuration");
+        eprintln!("  output.jsonl  - Output path for tagged corpus");
+        process::exit(1);
+    }
+
+    let corpus_path = &args[1];
+    let config_path = &args[2];
+    let output_path = &args[3];
+
+    if !Path::new(corpus_path).exists() {
+        eprintln!("Error: Corpus file not found: {}", corpus_path);
+        process::exit(1);
+    }
+
+    if !Path::new(config_path).exists() {
+        eprintln!("Error: Config file not found: {}", config_path);
+        process::exit(1);
+    }
+
+    let config_data = match std::fs::read_to_string(config_path) {
+        Ok(data) => data,
+        Err(e) => {
+            eprintln!("Error: Failed to read config file: {}", e);
+            process::exit(1);
+        }
+    };
+
+    let config: TagConfig = match serde_json::from_str(&config_data) {
+        Ok(cfg) => cfg,
+        Err(e) => {
+            eprintln!("Error: Failed to parse config: {}", e);
+            process::exit(1);
+        }
+    };
+
+    if let Err(e) = tag_corpus(corpus_path, &config, output_path) {
+        eprintln!("Error: Tagging failed: {}", e);
+        process::exit(1);
+    }
+}

data/ext/phrasekit/src/lib.rs

@@ -0,0 +1,104 @@
+mod manifest;
+mod matcher;
+mod payload;
+mod policy;
+
+use magnus::{define_module, function, method, prelude::*, Error, RArray, RHash, Ruby};
+use matcher::{Matcher as RustMatcher, Stats};
+use parking_lot::RwLock;
+use policy::MatchPolicy;
+use std::sync::Arc;
+
+type SharedMatcher = Arc<RwLock<Option<Arc<RustMatcher>>>>;
+
+#[magnus::wrap(class = "PhraseKit::NativeMatcher", free_immediately, size)]
+struct MatcherWrapper {
+    matcher: SharedMatcher,
+}
+
+impl MatcherWrapper {
+    fn new() -> Self {
+        Self {
+            matcher: Arc::new(RwLock::new(None)),
+        }
+    }
+
+    fn load(&self, automaton_path: String, payloads_path: String, manifest_path: String) -> Result<(), Error> {
+        let matcher = RustMatcher::load(&automaton_path, &payloads_path, &manifest_path)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load matcher: {}", e)))?;
+
+        let mut guard = self.matcher.write();
+        *guard = Some(Arc::new(matcher));
+
+        Ok(())
+    }
+
+    fn match_tokens(&self, token_ids: Vec<u32>, policy: String, max: usize) -> Result<RArray, Error> {
+        let guard = self.matcher.read();
+        let matcher = guard
+            .as_ref()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Matcher not loaded"))?;
+
+        let match_policy = MatchPolicy::from_str(&policy)
+            .ok_or_else(|| Error::new(magnus::exception::arg_error(), format!("Invalid policy: {}", policy)))?;
+
+        let matches = matcher.match_tokens(&token_ids, match_policy, max);
+
+        let result = RArray::new();
+        for m in matches {
+            let hash = RHash::new();
+            hash.aset("start", m.start)?;
+            hash.aset("end", m.end)?;
+            hash.aset("phrase_id", m.payload.phrase_id)?;
+            hash.aset("salience", m.payload.salience)?;
+            hash.aset("count", m.payload.count)?;
+            hash.aset("n", m.payload.n)?;
+            result.push(hash)?;
+        }
+
+        Ok(result)
+    }
+
+    fn stats(&self) -> Result<RHash, Error> {
+        let guard = self.matcher.read();
+        let matcher = guard
+            .as_ref()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Matcher not loaded"))?;
+
+        let stats = Stats::from_matcher(matcher);
+        let hash = RHash::new();
+
+        hash.aset("version", stats.version)?;
+        hash.aset("loaded_at", stats.loaded_at.duration_since(std::time::UNIX_EPOCH).unwrap().as_millis() as u64)?;
+        hash.aset("num_patterns", stats.num_patterns)?;
+        hash.aset("heap_mb", stats.heap_mb)?;
+        hash.aset("hits_total", stats.hits_total)?;
+        hash.aset("p50_us", stats.p50_us)?;
+        hash.aset("p95_us", stats.p95_us)?;
+        hash.aset("p99_us", stats.p99_us)?;
+
+        Ok(hash)
+    }
+
+    fn healthcheck(&self) -> Result<bool, Error> {
+        let guard = self.matcher.read();
+        guard
+            .as_ref()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Matcher not loaded"))?;
+        Ok(true)
+    }
+}
+
+#[magnus::init]
+fn init(ruby: &Ruby) -> Result<(), Error> {
+    let module = define_module("PhraseKit")?;
+    let class = module.define_class("NativeMatcher", ruby.class_object())?;
+
+    class.define_singleton_method("new", function!(MatcherWrapper::new, 0))?;
+    class.define_method("load", method!(MatcherWrapper::load, 3))?;
+    class.define_method("match_tokens", method!(MatcherWrapper::match_tokens, 3))?;
+    class.define_method("stats", method!(MatcherWrapper::stats, 0))?;
+    class.define_method("healthcheck", method!(MatcherWrapper::healthcheck, 0))?;
+
+    Ok(())
+}

data/ext/phrasekit/src/manifest.rs

@@ -0,0 +1,88 @@
+use serde::{Deserialize, Serialize};
+use std::fs::File;
+use std::io::BufReader;
+use std::path::Path;
+use thiserror::Error;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Manifest {
+    pub version: String,
+    pub tokenizer: String,
+    pub num_patterns: usize,
+    pub min_count: Option<u32>,
+    pub salience_threshold: Option<f32>,
+    pub built_at: String,
+    pub separator_id: u32,
+}
+
+#[derive(Error, Debug)]
+pub enum ManifestError {
+    #[error("IO error: {0}")]
+    Io(#[from] std::io::Error),
+
+    #[error("JSON parse error: {0}")]
+    Json(#[from] serde_json::Error),
+
+    #[error("Invalid manifest: {0}")]
+    #[allow(dead_code)]
+    Invalid(String),
+}
+
+impl Manifest {
+    pub fn load<P: AsRef<Path>>(path: P) -> Result<Self, ManifestError> {
+        let file = File::open(path)?;
+        let reader = BufReader::new(file);
+        let manifest: Manifest = serde_json::from_reader(reader)?;
+
+        if manifest.separator_id == 0 {
+            return Err(ManifestError::Invalid(
+                "separator_id must be non-zero".to_string(),
+            ));
+        }
+
+        Ok(manifest)
+    }
+
+    #[allow(dead_code)]
+    pub fn validate_compatible(&self, other: &Manifest) -> Result<(), ManifestError> {
+        if self.tokenizer != other.tokenizer {
+            return Err(ManifestError::Invalid(format!(
+                "Tokenizer mismatch: expected {}, got {}",
+                self.tokenizer, other.tokenizer
+            )));
+        }
+
+        if self.separator_id != other.separator_id {
+            return Err(ManifestError::Invalid(format!(
+                "Separator ID mismatch: expected {}, got {}",
+                self.separator_id, other.separator_id
+            )));
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_manifest_deserialize() {
+        let json = r#"{
+            "version": "pk-2025-09-25-01",
+            "tokenizer": "scientist-v1",
+            "num_patterns": 1287345,
+            "min_count": 20,
+            "salience_threshold": 1.0,
+            "built_at": "2025-09-25T18:44:00Z",
+            "separator_id": 4294967294
+        }"#;
+
+        let manifest: Manifest = serde_json::from_str(json).unwrap();
+        assert_eq!(manifest.version, "pk-2025-09-25-01");
+        assert_eq!(manifest.tokenizer, "scientist-v1");
+        assert_eq!(manifest.num_patterns, 1287345);
+        assert_eq!(manifest.separator_id, 4294967294);
+    }
+}