Merge pull request #124 from oscar-project/dev-checksum

Optional corpus checksum + have language folders rather than flat files

Cargo.toml

@@ -37,6 +37,8 @@ twox-hash = "1.6"
 glob = "0.3.0"
 sha2 = "0.9.5"
 
+oscar-tools = { version = "0.4.0", optional = true }
+
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 schemars = "0.8.3"
@@ -59,6 +61,7 @@ ctclib-pp = { version = "0.2.0", optional = true }
 
 [features]
 kenlm = ["dep:ctclib-pp"]
+checksum = ["dep:oscar-tools"]
 
 [dev-dependencies]
 rand_distr = "0.4.2"
@@ -97,4 +100,9 @@ ci = ["github"]
 # The installers to generate for each app
 installers = []
 # Target platforms to build apps for (Rust target-triple syntax)
-targets = ["x86_64-unknown-linux-gnu", "x86_64-apple-darwin", "x86_64-pc-windows-msvc", "aarch64-apple-darwin"]
+targets = [
+ "x86_64-unknown-linux-gnu",
+ "x86_64-apple-darwin",
+ "x86_64-pc-windows-msvc",
+ "aarch64-apple-darwin",
+]

src/cli.rs

@@ -195,4 +195,8 @@ pub struct Pipeline {
 
  #[structopt(short = "c", long = "comp", help = "Enables zstd compression")]
  pub comp: bool,
+
+ #[cfg(feature = "checksum")]
+ #[structopt(long = "checksum", help = "compute checksums")]
+ pub checksum: bool,
 }

src/io/langfiles.rs

@@ -99,7 +99,13 @@ impl LangFilesDoc {
  } else {
  None
  };
- let w = Writer::new(dst, lang, part_size_bytes, comp)?;
+
+ // add lang subfolder
+ let mut subfolder = dst.to_path_buf();
+ subfolder.push(lang.to_string());
+ std::fs::create_dir(&subfolder)?;
+
+ let w = Writer::new(&subfolder, lang, part_size_bytes, comp)?;
 
  Ok(Arc::new(Mutex::new(w)))
  }
@@ -219,7 +225,7 @@ mod tests {
  w.flush().unwrap();
  }
  let mut read_path = PathBuf::from(dst.path());
- read_path.push("en.jsonl");
+ read_path.push("en/en.jsonl");
 
  let b = File::open(read_path).unwrap();
  let doc_from_file: Document = serde_json::from_reader(b).unwrap();
@@ -252,7 +258,7 @@ mod tests {
  }
 
  let mut read_path = PathBuf::from(dst.path());
- read_path.push("en.jsonl.zstd");
+ read_path.push("en/en.jsonl.zstd");
 
  let b = File::open(&read_path).unwrap();
  let dec = zstd::decode_all(b).unwrap();

src/main.rs

@@ -73,6 +73,8 @@ async fn main() -> Result<(), error::Error> {
  p.kenlms_path,
  p.split.map(|size_mbytes| size_mbytes * 1_000_000),
  p.comp,
+ #[cfg(feature = "checksum")]
+ p.checksum,
  );
  p.run()?;
 

src/pipelines/oscardoc/pipeline.rs

@@ -60,8 +60,15 @@ pub struct OscarDoc {
  kenlms_path: Option<PathBuf>,
  split: Option<u64>, // in bytes
  comp: bool,
+ #[cfg(feature = "checksum")]
+ checksum: bool,
 }
 
+#[cfg(feature = "checksum")]
+use oscar_tools::Checksum;
+#[cfg(feature = "checksum")]
+impl Checksum for OscarDoc {}
+
 impl OscarDoc {
  pub fn new(
  src: PathBuf,
@@ -71,6 +78,7 @@ impl OscarDoc {
  kenlms_path: Option<PathBuf>,
  split: Option<u64>,
  comp: bool,
+ #[cfg(feature = "checksum")] checksum: bool,
  ) -> Self {
  if blocklist.is_none() {
  warn!("No blocklist folder specified! No adult content tagging will be done.");
@@ -85,6 +93,8 @@ impl OscarDoc {
  kenlms_path,
  split,
  comp,
+ #[cfg(feature = "checksum")]
+ checksum,
  }
  }
 
@@ -519,6 +529,16 @@ impl Pipeline<()> for OscarDoc {
  // flush writers
  info!("Flushing writers");
  langfiles.flush_all()?;
+ std::mem::drop(langfiles);
+ //compute checksums
+ #[cfg(feature = "checksum")]
+ {
+ if self.checksum {
+ info!("Checksumming");
+ Self::checksum_folder(&self.dst, 1).unwrap();
+ }
+ }
+
  info!("Done");
  Ok(())
  }

src/pipelines/oscardoc/types/location.rs

@@ -191,8 +191,6 @@ impl Default for Location {
 
 #[cfg(test)]
 mod tests {
-
-
 
  use super::Location;
  use super::LocationBuilder;

src/transformers/mod.rs

@@ -24,13 +24,13 @@ pub use annotate::Annotate;
 pub use annotate::Annotator;
 pub use content_detector::ContentDetector;
 pub use header::Header;
-pub use lsh::LSH;
 #[cfg(feature = "kenlm")]
 pub use kenlm::AdultDetector;
 #[cfg(feature = "kenlm")]
 pub use kenlm::AdultDetectorBuilder;
 #[cfg(feature = "kenlm")]
 pub use kenlm::Models;
+pub use lsh::LSH;
 pub use noisy::Noisy;
 pub use sentence_filter::Conv;
 pub use sentence_filter::RemoveShortSentences;

tests/oscardoc_rebuild.rs

@@ -24,7 +24,12 @@ fn gen_corpus() {
  let kenlm = Path::new("res/kenlm/").to_path_buf();
 
  //TODO test with custom blocklists
+ #[cfg(feature = "checksum")]
+ let pipeline = OscarDoc::new(src, dst, lid, Some(bl), Some(kenlm), None, false, false);
+
+ #[cfg(not(feature = "checksum"))]
  let pipeline = OscarDoc::new(src, dst, lid, Some(bl), Some(kenlm), None, false);
+
  pipeline.run().expect(
  "Ensure to have shards in res/shards, lid.176.bin at root and blocklist at res/blocklist",
  );