Add dataset package

1 files changed, 224 insertions, 0 deletions

diff --git a/dataset/download.go b/dataset/download.go
new file mode 100644
index 0000000..0df6ebe
--- /dev/null
+++ b/dataset/download.go
@@ -0,0 +1,224 @@
+package dataset
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"iter"
+	"net/http"
+	"net/url"
+	"os"
+	"path/filepath"
+	"slices"
+)
+
+const Base = "https://datasets-server.huggingface.co/parquet"
+
+func Download(dataset, config, root string) error {
+	var pl *parquetListing
+
+	if l, err := listing(dataset, config); err != nil {
+		return err
+	} else {
+		pl = l
+	}
+
+	splits := pl.Splits()
+
+	for _, split := range splits {
+		dir := filepath.Join(root, split)
+
+		if err := os.MkdirAll(dir, 0o755); err != nil {
+			return err
+		}
+
+		total := int64(0)
+
+		for pf := range pl.Split(split) {
+			name := filepath.Join(dir, filepath.Base(pf.Filename))
+
+			if n, err := download(pf, name); err != nil {
+				return err
+			} else {
+				total += n
+
+				fmt.Printf("[%s] downloaded %s (%d bytes), total=%d\n", split, filepath.Base(name), n, total)
+			}
+		}
+	}
+
+	return nil
+}
+
+type parquetListing struct {
+	ParquetFiles []parquetFile `json:"parquet_files"`
+}
+
+type parquetFile struct {
+	Dataset  string `json:"dataset"`
+	Config   string `json:"config"`
+	Split    string `json:"split"`
+	URL      string `json:"url"`
+	Filename string `json:"filename"`
+	Size     int64  `json:"size"`
+}
+
+func (pl *parquetListing) Splits() []string {
+	seen := make(map[string]struct{})
+
+	for _, ps := range pl.ParquetFiles {
+		if ps.Split == "" {
+			continue // TODO possible?
+		}
+
+		seen[ps.Split] = struct{}{}
+	}
+
+	splits := make([]string, 0)
+
+	for k := range seen {
+		splits = append(splits, k)
+	}
+
+	slices.Sort(splits)
+
+	return splits
+}
+
+func (pl *parquetListing) Split(split string) iter.Seq[parquetFile] {
+	return func(yield func(pf parquetFile) bool) {
+		for _, pf := range pl.ParquetFiles {
+			if pf.Split != split {
+				continue
+			}
+
+			if !yield(pf) {
+				return
+			}
+		}
+	}
+}
+
+func listing(dataset, config string) (*parquetListing, error) {
+	var parquet *url.URL
+
+	if u, err := url.Parse(Base); err != nil {
+		return nil, err
+	} else {
+		parquet = u
+	}
+
+	query := parquet.Query()
+
+	query.Set("dataset", dataset)
+	query.Set("config", config)
+
+	parquet.RawQuery = query.Encode()
+
+	var req *http.Request
+
+	if r, err := http.NewRequest("GET", parquet.String(), nil); err != nil {
+		return nil, err
+	} else {
+		req = r
+	}
+
+	var resp *http.Response
+
+	if r, err := http.DefaultClient.Do(req); err != nil {
+		return nil, err
+	} else {
+		resp = r
+
+		defer resp.Body.Close()
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, errors.New("http not ok")
+	}
+
+	var pl parquetListing
+
+	if err := json.NewDecoder(resp.Body).Decode(&pl); err != nil {
+		return nil, err
+	}
+
+	return &pl, nil
+}
+
+func download(pf parquetFile, name string) (int64, error) {
+	if st, err := os.Stat(name); err == nil && st.Size() == pf.Size {
+		return pf.Size, nil
+	}
+
+	var req *http.Request
+
+	if r, err := http.NewRequest("GET", pf.URL, nil); err != nil {
+		return 0, err
+	} else {
+		req = r
+	}
+
+	var resp *http.Response
+
+	if r, err := http.DefaultClient.Do(req); err != nil {
+		return 0, err
+	} else {
+		resp = r
+
+		defer resp.Body.Close()
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return 0, errors.New("http not ok")
+	}
+
+	tmp := name + ".part"
+
+	var file *os.File
+
+	ok := false
+
+	if f, err := os.Create(tmp); err != nil {
+		return 0, err
+	} else {
+		file = f
+	}
+
+	defer func() {
+		_ = file.Close()
+
+		if !ok {
+			_ = os.Remove(tmp)
+		}
+	}()
+
+	written := int64(0)
+
+	if n, err := io.Copy(file, resp.Body); err != nil {
+		return 0, err
+	} else {
+		written = n
+	}
+
+	if err := file.Sync(); err != nil {
+		return 0, err
+	}
+
+	if pf.Size > 0 && written != pf.Size {
+		return 0, errors.New("size mismatch")
+	}
+
+	if err := file.Close(); err != nil {
+		return 0, err
+	}
+
+	if err := os.Rename(tmp, name); err != nil {
+		return 0, err
+	}
+
+	ok = true
+
+	return written, nil
+}