<commit_before># vim: set ts=2 et:
import argparse
import io
import sys
import threading
try:
  import zmq
except ImportError:
  zmq = None

from ..reader import Reader
from ..writer import Writer


def stream_coroutine(istream, ostream, doc_class=None, automagic=False):
  reader = Reader(istream, doc_class, automagic)
  writer = Writer(ostream, reader.doc_schema)
  for doc in reader:
    res = yield(doc)
    writer.write(res or doc)


def zmq_coroutine(context, dealer_url, doc_class=None, automagic=False):
  istream = io.BytesIO()
  ostream = io.BytesIO()
  reader = Reader(istream, doc_class, automagic)
  writer = Writer(ostream, reader.doc_schema)
  socket = context.socket(zmq.REP)
  socket.connect(dealer_url)
  while True:
    msg = socket.recv()
    istream.write(msg)
    istream.seek(0)
    for doc in reader:
      res = yield(doc)
      writer.write(res or doc)
    ostream.seek(0)
    socket.send(ostream.getvalue())
    istream.truncate(0)
    ostream.truncate(0)


arg_parser = argparse.ArgumentParser(add_help=False)
if zmq:
  _megroup = arg_parser.add_mutually_exclusive_group()
  _megroup.add_argument('--serve', dest='serve_url', metavar="ADDRESS", default=None, help='Serve from the specified address, e.g. tcp://*:7300')
  _megroup.add_argument('--worker', dest='worker_url', metavar="ADDRESS", default=None, help='Acquire work from the specified address')<commit_msg>Fix zmq processing<commit_after># vim: set ts=2 et:
import argparse
import io
import sys
import threading
try:
  import zmq
except ImportError:
  zmq = None

from ..reader import Reader
from ..writer import Writer


def stream_coroutine(istream, ostream, doc_class=None, automagic=False):
  reader = Reader(istream, doc_class, automagic)
  writer = Writer(ostream, reader.doc_schema)
  for doc in reader:
    res = yield(doc)
    writer.write(res or doc)


def zmq_coroutine(context, dealer_url, doc_class=None, automagic=False):
  # FIXME: reduce overhead of reader/writer creation
  ostream = io.BytesIO()
  socket = context.socket(zmq.REP)
  socket.connect(dealer_url)
  while True:
    msg = socket.recv()
    istream = io.BytesIO(msg)
    istream.seek(0)
    reader = Reader(istream, doc_class, automagic)
    writer = Writer(ostream, reader.doc_schema)
    for doc in reader:
      res = yield(doc)
      writer.write(res or doc)
    ostream.seek(0)
    socket.send(ostream.getvalue())
    ostream.truncate(0)


arg_parser = argparse.ArgumentParser(add_help=False)
if zmq:
  _megroup = arg_parser.add_mutually_exclusive_group()
  _megroup.add_argument('--serve', dest='serve_url', metavar="ADDRESS", default=None, help='Serve from the specified address, e.g. tcp://*:7300')
  _megroup.add_argument('--worker', dest='worker_url', metavar="ADDRESS", default=None, help='Acquire work from the specified address')
{'fraction_non_alphanumeric': 0.11708185053380783, 'fraction_numerical': 0.006049822064056939, 'mean_word_length': 4.5884691848906565, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '220307', 'n_tokens_mistral': 974, 'n_tokens_neox': 939, 'n_words': 233}