meta-pytorch / meta-pytorch/data
Open for contribution on utility nodes like `Filter`, `Shuffler`, `Header`, `Cycler`?
Open
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.3k
- Forks
- 179
- Avg merge
- 6d 1h
- Merged PRs (30d)
- 2
Description
Hi, do you think this kind of nodes would be in the scope of Torchdata? Then I'm down to open a PR to add them. with remaining and testing, for sure.
import logging
import random
from collections import deque
from typing import Any, Callable, Deque, Dict, Optional, TypeVar, Optional
from torchdata.nodes import BaseNode
logger = logging.getLogger(__name__)
X = TypeVar("X")
T = TypeVar("T")
U = TypeVar("U")
class Filter(BaseNode[T]):
"""Node that filters items from source node based on predicate function."""
SOURCE_KEY = "source"
def __init__(self, source_node: BaseNode[T], filter_fn: Callable[[T], bool]):
super().__init__()
self.source = source_node
self.filter_fn = filter_fn
def reset(self, initial_state: Optional[Dict[str, Any]] = None):
super().reset(initial_state)
self.source.reset(initial_state.get(self.SOURCE_KEY) if initial_state else None)
def next(self) -> T:
while True:
item = next(self.source)
if self.filter_fn(item):
return item
def get_state(self) -> Dict[str, Any]:
return {self.SOURCE_KEY: self.source.state_dict()}
class Shuffler(BaseNode[T]):
"""Node that shuffles items from source node using a buffer."""
SOURCE_KEY = "source"
def __init__(self, source_node: BaseNode[T], buffer_size: int, seed: Optional[int] = None):
super().__init__()
if buffer_size < 1:
raise ValueError("Buffer size must be at least 1")
self.source = source_node
self.buffer_size = buffer_size
self.buffer: Deque[T] = deque()
self.rng = random.Random(seed)
self._initial_seed = seed
def reset(self, initial_state: Optional[Dict[str, Any]] = None):
super().reset(initial_state)
self.buffer.clear()
if initial_state is not None:
self.source.reset(initial_state.get(self.SOURCE_KEY))
self.rng.setstate(initial_state["rng_state"])
else:
self.source.reset()
if self._initial_seed is not None:
self.rng = random.Random(self._initial_seed)
def _fill_buffer(self) -> bool:
"""Fill buffer with items from source. Returns True if any items were added."""
try:
while len(self.buffer) < self.buffer_size:
self.buffer.append(next(self.source))
return True
except StopIteration:
return len(self.buffer) > 0
def next(self) -> T:
if not self.buffer and not self._fill_buffer():
raise StopIteration
# Randomly select and remove an item from the buffer
idx = self.rng.randrange(len(self.buffer))
item = self.buffer[idx]
self.buffer[idx] = self.buffer[-1]
self.buffer.pop()
# Try to refill buffer
self._fill_buffer()
return item
def get_state(self) -> Dict[str, Any]:
return {self.SOURCE_KEY: self.source.state_dict(), "rng_state": self.rng.getstate()}
class Header(BaseNode[T]):
"""Node that yields only the first N items from source node."""
SOURCE_KEY = "source"
def __init__(self, source_node: BaseNode[T], n: int):
super().__init__()
if n < 0:
raise ValueError("n must be non-negative")
self.source = source_node
self.n = n
self._count = 0
def reset(self, initial_state: Optional[Dict[str, Any]] = None):
super().reset(initial_state)
self.source.reset(initial_state.get(self.SOURCE_KEY) if initial_state else None)
if initial_state is not None:
self._count = initial_state["count"]
else:
self._count = 0
def next(self) -> T:
if self._count >= self.n:
raise StopIteration
item = next(self.source)
self._count += 1
return item
def get_state(self) -> Dict[str, Any]:
return {self.SOURCE_KEY: self.source.state_dict(), "count": self._count}
class Cycler(BaseNode[T]):
"""Node that cycles through source node indefinitely."""
SOURCE_KEY = "source"
def __init__(self, source_node: BaseNode[T]):
super().__init__()
self.source = source_node
self._cycle_count: int = 0
def reset(self, initial_state: Optional[Dict[str, Any]] = None):
super().reset(initial_state)
if initial_state is not None:
self._cycle_count = initial_state["cycle_count"]
self.source.reset(initial_state.get(self.SOURCE_KEY))
else:
self._cycle_count = 0
self.source.reset(None)
def next(self) -> T:
try:
return next(self.source)
except StopIteration:
self._cycle_count += 1
self.source.reset(None)
return next(self.source)
def get_state(self) -> Dict[str, Any]:
return {self.SOURCE_KEY: self.source.state_dict(), "cycle_count": self._cycle_count}
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reading torchdata.nodes.BaseNode and the proposed Filter, Shuffler, Header, and Cycler implementations in the issue. Check how existing nodes handle reset, state, iteration, and testing; the issue is done only after maintainers define the supported utility-node scope and corresponding behavior and tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- data
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100