Backup script stream_loader.py

scripts/stream_loader.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+"""
+Custom streaming data loader for AGILLM training
+Pulls from stream_server on scraper box via HTTP
+Drop-in replacement for HuggingFace dataset streaming
+"""
+import requests
+import json
+from typing import Iterator, Dict, Any
+
+class ScraperStreamDataset:
+    """
+    Streams training data from the scraper server.
+    Compatible with AGILLM's _stream() interface.
+    """
+    def __init__(
+        self,
+        server_url: str = "http://localhost:8888",  # Will be SSH tunneled
+        batch_size: int = 100,
+        text_field: str = "text",
+        shuffle: bool = True
+    ):
+        self.server_url = server_url
+        self.batch_size = batch_size
+        self.text_field = text_field
+        self.shuffle = shuffle
+        self._buffer = []
+    
+    def __iter__(self) -> Iterator[Dict[str, Any]]:
+        return self
+    
+    def __next__(self) -> Dict[str, Any]:
+        if not self._buffer:
+            self._fetch_batch()
+        if not self._buffer:
+            raise StopIteration
+        return self._buffer.pop(0)
+    
+    def _fetch_batch(self):
+        """Fetch a batch from stream server"""
+        endpoint = "/stream" if self.shuffle else "/sequential"
+        try:
+            resp = requests.get(
+                f"{self.server_url}{endpoint}",
+                params={"batch": self.batch_size},
+                stream=True,
+                timeout=30
+            )
+            for line in resp.iter_lines():
+                if line:
+                    try:
+                        obj = json.loads(line.decode('utf-8'))
+                        # Return in format trainer expects
+                        self._buffer.append({self.text_field: obj.get("text", "")})
+                    except json.JSONDecodeError:
+                        continue
+        except requests.RequestException as e:
+            print(f"[StreamLoader] Fetch error: {e}")
+    
+    def get_status(self) -> dict:
+        """Get server status"""
+        try:
+            resp = requests.get(f"{self.server_url}/status", timeout=10)
+            return resp.json()
+        except:
+            return {"error": "unreachable"}
+
+
+def create_stream_iterator(server_url: str = "http://localhost:8888", seed: int = 42):
+    """
+    Create iterator compatible with AGILLM's _stream() function.
+    Returns infinite iterator of {"text": "..."} dicts.
+    """
+    dataset = ScraperStreamDataset(server_url=server_url)
+    while True:
+        try:
+            yield next(dataset)
+        except StopIteration:
+            # Refill and continue
+            dataset._fetch_batch()
+            if dataset._buffer:
+                yield dataset._buffer.pop(0)
+
+
+# For testing
+if __name__ == "__main__":
+    import sys
+    url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:8888"
+    print(f"Testing stream from {url}")
+    
+    ds = ScraperStreamDataset(server_url=url, batch_size=5)
+    print(f"Status: {ds.get_status()}")
+    
+    for i, item in enumerate(ds):
+        text = item["text"]
+        print(f"Sample {i}: {len(text)} chars - {text[:100]}...")
+        if i >= 4:
+            break