datasets

package
v0.3.5 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Apr 16, 2026 License: Apache-2.0 Imports: 20 Imported by: 0

Documentation

Overview

Package datasets provides an interface for interacting with the HuggingFace datasets server and downloading dataset files, functioning similarly to the `hub` package.

Index

Constants

This section is empty.

Variables

View Source
var Acronyms = map[string]struct{}{
	"id":    {},
	"html":  {},
	"url":   {},
	"http":  {},
	"api":   {},
	"ascii": {},
	"css":   {},
	"dns":   {},
	"https": {},
	"json":  {},
	"ip":    {},
	"uri":   {},
	"tcp":   {},
}

Acronyms is a map of common abbreviations that should be fully capitalized when generating Go structs. For example, "id" becomes "ID" instead of "Id".

Functions

func CreateParquetReader

func CreateParquetReader[T any](ds *Dataset, config, split string) (*parquet.GenericReader[T], error)

CreateParquetReader creates a parquet reader for the given dataset, config and split. This is more flexible than an iterator because it allows for random reads.

It groups together all the files for the given config and split and creates a single reader for them. It will fix the schema to match the given type T.

func GenerateGoStructFromParquet

func GenerateGoStructFromParquet(parquetFilePath, rootStructName string) (string, error)

GenerateGoStructFromParquet reads the schema of a local Parquet file and generates Go source code defining the structures to hold its records.

func IterParquetFromDataset

func IterParquetFromDataset[T any](ds *Dataset, config, split string) iter.Seq2[T, error]

IterParquetFromDataset downloads all Parquet files associated with the dataset's config and split and iterates over all their records sequentially. It will yield an error and stop if there's an issue acquiring or reading the files.

func IterParquetFromDatasetAt

func IterParquetFromDatasetAt[T any](ds *Dataset, config, split string, at int64) iter.Seq2[T, error]

IterParquetFromDatasetAt provides the same functionality as IterParquetFromDataset but starts at the given record.

It uses CreateParquetReader to seek to the required position. It will yield an error and stop if there's an issue acquiring or reading the files.

func IterParquetFromFile

func IterParquetFromFile[T any](filePath string) iter.Seq2[T, error]

IterParquetFromFile iterates over the records of a local Parquet file. Yields the records of type T mapped from the underlying parquet columns.

func ParquetFixListSchema

func ParquetFixListSchema[T any](filePath string) (*parquet.Schema, error)

ParquetFixListSchema recursively transforms a struct-based schema to match a file's naming.

Types

type ConfigInfo

type ConfigInfo struct {
	Description  string             `json:"description"`
	Citation     string             `json:"citation"`
	Homepage     string             `json:"homepage"`
	License      string             `json:"license"`
	Features     map[string]Feature `json:"features"`
	BuilderName  string             `json:"builder_name"`
	ConfigName   string             `json:"config_name"`
	Version      map[string]any     `json:"version"`
	Splits       map[string]Split   `json:"splits"`
	DownloadSize int64              `json:"download_size"`
	DatasetSize  int64              `json:"dataset_size"`
}

ConfigInfo holds the info for a specific configuration of a dataset.

func (*ConfigInfo) String

func (c *ConfigInfo) String() string

String implements fmt.Stringer to pretty-print the ConfigInfo.

type Dataset

type Dataset struct {
	*hub.Repo

	// Files slice contains a list of backing ParquetFile entries for the dataset.
	// It is lazily populated via GetParquetFiles().
	Files []ParquetFile
	// contains filtered or unexported fields
}

Dataset from which one wants to download files or get info. It embeds a *hub.Repo to reuse its downloading, caching, and file URL resolving functionality.

func New

func New(id string) *Dataset

New creates a reference to a HuggingFace dataset given its id. The id typically includes owner/dataset name. E.g.: "microsoft/ms_marco".

func (*Dataset) Download

func (d *Dataset) Download(files ...ParquetFile) ([]string, error)

Download downloads the given parquet files. It returns the list of absolute paths where the files were downloaded.

func (*Dataset) DownloadCtx

func (d *Dataset) DownloadCtx(ctx context.Context, parquetFiles ...ParquetFile) (downloadedPaths []string, err error)

DownloadCtx is like Download but accepts a context for cancellation support.

func (*Dataset) DownloadDatasetInfo

func (d *Dataset) DownloadDatasetInfo(ctx context.Context, forceDownload bool) error

DownloadDatasetInfo downloads the dataset info using the datasets-server.

If forceDownload is set to true, it ignores the cached one.

func (*Dataset) DownloadParquetFilesInfo

func (d *Dataset) DownloadParquetFilesInfo(ctx context.Context, forceDownload bool) error

DownloadParquetFilesInfo downloads the parquet info using the datasets-server.

If forceDownload is set to true, it ignores the cached one.

func (*Dataset) GenerateGoStruct

func (d *Dataset) GenerateGoStruct(config, split string) (string, error)

GenerateGoStruct lists files for the given config and split, downloads one parquet file to inspect its schema, and returns Go source code defining the structures to hold its records.

func (*Dataset) GetParquetFiles

func (d *Dataset) GetParquetFiles() ([]ParquetFile, error)

GetParquetFiles returns the list of parquet files for the dataset. If it hasn't been downloaded or loaded from the cache yet, it loads it first.

func (*Dataset) Info

func (d *Dataset) Info() (*Info, error)

Info returns the Info structure about the dataset. If it hasn't been downloaded or loaded from the cache yet, it loads it first. It may return nil if there was an issue with the downloading of the Info json from HuggingFace. Try DownloadDatasetInfo to get an error.

func (*Dataset) ListFiles

func (d *Dataset) ListFiles(config, split string) ([]ParquetFile, error)

ListFiles returns the parquet files that match the given config and split. If config or split are empty, they are not used for filtering.

They are returned in the order they are listed in the dataset info.

func (*Dataset) String

func (d *Dataset) String() string

String returns a formatted, human-readable summary of the dataset.

func (*Dataset) WithAuth

func (d *Dataset) WithAuth(authToken string) *Dataset

WithAuth sets the authentication token to use during downloads and info requests. Setting it to empty ("") is the same as resetting and not using authentication.

func (*Dataset) WithCacheDir

func (d *Dataset) WithCacheDir(cacheDir string) *Dataset

WithCacheDir sets the cache directory for downloaded files to the given directory.

func (*Dataset) WithEndpoint

func (d *Dataset) WithEndpoint(endpoint string) *Dataset

WithEndpoint sets the HuggingFace endpoint to use. Default is "https://huggingface.co" or, if set, the environment variable HF_ENDPOINT.

func (*Dataset) WithRevision

func (d *Dataset) WithRevision(revision string) *Dataset

WithRevision sets the revision to use for this Dataset, defaults to "main", but can be set to a commit-hash value.

type Feature

type Feature struct {
	Type  string `json:"_type,omitempty"`
	DType string `json:"dtype,omitempty"`

	// SubFeature is used when Type is "Sequence" or "Array". It will be a map[string]Feature.
	// If the feature is a single value, it will be stored under the empty string key ("").
	SubFeature map[string]Feature `json:"-"`

	// Used by features like ClassLabel
	Names []string `json:"names,omitempty"`
}

Feature represents a dataset feature definition (column schema).

func (Feature) MarshalJSON

func (f Feature) MarshalJSON() ([]byte, error)

MarshalJSON implements custom JSON marshaling to flatten single SubFeatures.

func (*Feature) UnmarshalJSON

func (f *Feature) UnmarshalJSON(b []byte) error

UnmarshalJSON implements custom JSON unmarshaling to handle dynamic "feature" values.

type FieldDef

type FieldDef struct {
	GoName   string
	GoType   string
	JSONName string
	IsList   bool
}

type Info

type Info struct {
	// DatasetInfo is keyed on the "config" name.
	// Some datasets use "config" names to denote versions, e.g. "v1.1".
	DatasetInfo map[string]ConfigInfo `json:"dataset_info"`
}

Info holds information about a dataset returned by the datasets-server. See the API at: https://datasets-server.huggingface.co/info?dataset=<id>

func (*Info) String

func (i *Info) String() string

String implements fmt.Stringer to pretty-print the Info.

type ParquetFile

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"`
}

type ParquetFileResponse

type ParquetFileResponse struct {
	ParquetFiles []ParquetFile `json:"parquet_files"`
}

type Split

type Split struct {
	DatasetName string `json:"dataset_name,omitempty"`
	Name        string `json:"name"`
	NumBytes    int64  `json:"num_bytes"`
	NumExamples int64  `json:"num_examples"`
}

Split holds information about a dataset split (e.g. "train", "test", "validation").

type StructDef

type StructDef struct {
	Name   string
	Fields []FieldDef
}

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL