Add sample

Author: pierotofy

README.md

@@ -18,6 +18,16 @@ print(src, tgt)
 # source.txt.shuffled target.txt.shuffled
 ```
 
+You can also simultaneously sample/isolate a certain number of sentences from the dataset (which are then removed from the shuffled result)
+
+```
+from fastshuffle import file_shuffle_sample
+
+src, tgt, src_sample, tgt_sample = file_shuffle("source.txt", "target.txt", 5) # Sample 5 sentences
+print(src, tgt, src_sample, tgt_sample)
+# source.txt.shuffled target.txt.shuffled source.txt.shuffled.sample target.txt.shuffled.sample
+```
+
 ## Notes
 
 Source and target must have the same number of lines. No validation checks are made.

fastshuffle.cpp

@@ -8,6 +8,8 @@ namespace py = pybind11;
 
 PYBIND11_MODULE(fastshuffle, m) {
     m.def("file_shuffle", &shuffle);
+    m.def("file_shuffle_sample", &shuffle_sample);
+
 
 #ifdef VERSION_INFO
     m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO);

main.cpp

@@ -7,6 +7,7 @@ int main(int argc, char **argv) {
     options.add_options()
         ("s,src", "Input source corpus", cxxopts::value<std::string>())
         ("t,tgt", "Input target corpus", cxxopts::value<std::string>())
+        ("sample", "Also sample these many sentences", cxxopts::value<long long>()->default_value("0"))
         ("h,help", "Print usage")
         ;
     options.parse_positional({ "src", "tgt" });
@@ -30,10 +31,19 @@ int main(int argc, char **argv) {
     try {
         const auto source = result["src"].as<std::string>();
         const auto target = result["tgt"].as<std::string>();
+        const auto sample = result["sample"].as<long long>();
 
-        auto result = shuffle(source, target);
-        std::cout << "W\t" << std::get<0>(result) << std::endl;
-        std::cout << "W\t" << std::get<1>(result) << std::endl;
+        if (sample > 0){
+            auto result = shuffle_sample(source, target, sample);
+            std::cout << "W\t" << std::get<0>(result) << std::endl;
+            std::cout << "W\t" << std::get<1>(result) << std::endl;
+            std::cout << "W\t" << std::get<2>(result) << std::endl;
+            std::cout << "W\t" << std::get<3>(result) << std::endl;
+        }else{
+            auto result = shuffle(source, target);
+            std::cout << "W\t" << std::get<0>(result) << std::endl;
+            std::cout << "W\t" << std::get<1>(result) << std::endl;
+        }
     }catch (std::exception &e) {
         std::cerr << "Error: " << e.what() << std::endl;
         return EXIT_FAILURE;

setup.py

@@ -124,7 +124,7 @@ def build_extension(self, ext: CMakeExtension) -> None:
 
 setup(
     name="fastshuffle",
-    version="1.0.0",
+    version="1.0.1",
     author="Piero Toffanin",
     author_email="[email protected]",
     url="https://github.com/LibreTranslate/FastShuffle",

shuffle.cpp

@@ -71,12 +71,20 @@ size_t line_count(mmap_handle_t mht){
 }
 
 std::tuple<std::string, std::string> shuffle(const std::string &src, const std::string &tgt){
+    auto res = shuffle_sample(src, tgt, 0);
+    return std::make_tuple(std::get<0>(res), std::get<1>(res));
+}
+
+std::tuple<std::string, std::string, std::string, std::string> shuffle_sample(const std::string &src, const std::string &tgt, long long sample){
     mmap_handle_t smht = mmap_open(src);
     mmap_handle_t tmht = mmap_open(tgt);
 
     std::string src_out = src + ".shuffled";
     std::string tgt_out = tgt + ".shuffled";
 
+    std::string src_sample_out = src_out + ".sample";
+    std::string tgt_sample_out = tgt_out + ".sample";
+
     // size_t src_count = line_count(smht);
 
     std::vector<line_off_t> offsets;
@@ -122,12 +130,38 @@ std::tuple<std::string, std::string> shuffle(const std::string &src, const std::
     std::ofstream src_of(src_out, std::ios::trunc);
     std::ofstream tgt_of(tgt_out, std::ios::trunc);
 
+    std::ofstream *src_sample_of = nullptr;
+    std::ofstream *tgt_sample_of = nullptr;
+    
+    if (sample > 0){
+        src_sample_of = new std::ofstream(src_sample_out, std::ios::trunc);
+        tgt_sample_of = new std::ofstream(tgt_sample_out, std::ios::trunc);
+    }
+
     // std::cout << offsets.size() << std::endl << std::endl;
     for (size_t i = 0; i < offsets.size(); i++){
-        src_of.write(offsets[i].src_start, offsets[i].src_end - offsets[i].src_start);
-        src_of.write("\n", 1);
-        tgt_of.write(offsets[i].tgt_start, offsets[i].tgt_end - offsets[i].tgt_start);
-        tgt_of.write("\n", 1);
+        if (sample > 0 && i < sample){
+            src_sample_of->write(offsets[i].src_start, offsets[i].src_end - offsets[i].src_start);
+            src_sample_of->write("\n", 1);
+            tgt_sample_of->write(offsets[i].tgt_start, offsets[i].tgt_end - offsets[i].tgt_start);
+            tgt_sample_of->write("\n", 1);
+        }else{
+            src_of.write(offsets[i].src_start, offsets[i].src_end - offsets[i].src_start);
+            src_of.write("\n", 1);
+            tgt_of.write(offsets[i].tgt_start, offsets[i].tgt_end - offsets[i].tgt_start);
+            tgt_of.write("\n", 1);
+        }
+    }
+
+    if (tgt_sample_of != nullptr){ 
+        tgt_sample_of->close();
+        delete tgt_sample_of;
+        tgt_sample_of = nullptr;
+    }
+    if (src_sample_of != nullptr){ 
+        src_sample_of->close();
+        delete src_sample_of;
+        src_sample_of = nullptr;
     }
 
     src_of.close();
@@ -136,5 +170,5 @@ std::tuple<std::string, std::string> shuffle(const std::string &src, const std::
     mmap_close(smht);
     mmap_close(tmht);
 
-    return std::make_tuple(src_out, tgt_out);
+    return std::make_tuple(src_out, tgt_out, src_sample_out, tgt_sample_out);
 }

shuffle.hpp

@@ -42,4 +42,5 @@ typedef struct{
 mmap_handle_t mmap_open(const std::string &file);
 void mmap_close(mmap_handle_t mht);
 size_t line_count(mmap_handle_t mht);
-std::tuple<std::string, std::string> shuffle(const std::string &src, const std::string &tgt);
+std::tuple<std::string, std::string> shuffle(const std::string &src, const std::string &tgt);
+std::tuple<std::string, std::string, std::string, std::string> shuffle_sample(const std::string &src, const std::string &tgt, long long sample);

tests/test_shuffle.py

@@ -1,5 +1,5 @@
 import pytest
-from fastshuffle import file_shuffle
+from fastshuffle import file_shuffle, file_shuffle_sample
 import os
 
 def test_shuffle():
@@ -15,4 +15,17 @@ def test_shuffle():
     assert sum(1 for _ in open(src)) == sum(1 for _ in open(in_src))
     assert sum(1 for _ in open(tgt)) == sum(1 for _ in open(in_tgt))
 
+
+def test_shuffle_sample():
+    cwd = os.path.dirname(__file__)
+
+    in_src = os.path.join(cwd, "data", "src.txt")
+    in_tgt = os.path.join(cwd, "data", "tgt.txt")
+    src, tgt, ssrc, stgt = file_shuffle_sample(in_src, in_tgt, 1)
+
+    assert os.path.isfile(src)
+    assert os.path.isfile(tgt)
+    assert os.path.isfile(ssrc)
+    assert os.path.isfile(stgt)
+