dev-tools/scripts/create_line_file_docs.py - lucene - Git at Google

 #!/usr/bin/env python3
 # Licensed to the Apache Software Foundation (ASF) under one or more
 # contributor license agreements.  See the NOTICE file distributed with
 # this work for additional information regarding copyright ownership.
 # The ASF licenses this file to You under the Apache License, Version 2.0
 # (the "License"); you may not use this file except in compliance with
 # the License.  You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.

 import gzip
 import os
 import random
 import re
 import shutil
 import subprocess
 import tempfile
 import time
 import urllib.request
 from io import TextIOWrapper

 DEBUG = False

 TARGET_DOC_CHARS = 1024


 def compress_with_seek_points(file_name_in: str, file_name_out: str, num_seek_points: int):
   bytes_per_chunk = os.path.getsize(file_name_in) / num_seek_points

   seek_points: list[int] = []

   if os.path.exists(file_name_out):
     os.remove(file_name_out)

   with open(file_name_in, "rb") as f_in:
     f_out = None

     bytes_in_chunk = 0

     chunk_count = 0

     while True:
       if f_out is None:
         if os.path.exists(file_name_out):
           seek_points.append(os.path.getsize(file_name_out))
           print("  create chunk %s at pos=%s" % (chunk_count, seek_points[-1]))
         else:
           print("  create chunk %s at pos=0" % chunk_count)
         f_out = gzip.open(file_name_out, "ab")
         chunk_count += 1

       line = f_in.readline()
       if len(line) == 0:
         break

       bytes_in_chunk += len(line)
       f_out.write(line)

       if bytes_in_chunk > bytes_per_chunk and chunk_count < num_seek_points:
         f_out.close()
         f_out = None
         bytes_in_chunk = 0

   with open(file_name_out[:-3] + ".seek", "w") as f_out:
     for seek_point in seek_points:
       f_out.write("%d\n" % seek_point)


 re_tag = re.compile(r"<[^>]+?>")
 re_newlines = re.compile(r"\n+")
 re_space = re.compile(r"\s")

 # used to find word break, for splitting docs into ~1 KB sized smaller docs:
 re_next_non_word_character = re.compile(r"\W", re.UNICODE)

 EUROPARL_V7_URL = "https://www.statmt.org/europarl/v7/europarl.tgz"


 def split_docs(all_out: TextIOWrapper, title_string: str, date_string: str, body_string: str):
   """Splits docs into smallish (~1 KB) sized docs, repeating same title and date"""
   doc_count = 0
   while len(body_string) > 0:
     char_count = int(random.gauss(TARGET_DOC_CHARS, TARGET_DOC_CHARS / 4))
     if char_count < 64:
       # trimmed normal?
       continue

     m = re_next_non_word_character.search(body_string, char_count)
     if m is not None:
       char_count = m.start(0)
     else:
       char_count = len(body_string)

     body_string_fragment = body_string[:char_count].strip()

     # print('write title %d, body %d' % (len(title_string), len(body_string_fragment)))
     all_out.write("%s\t%s\t%s\n" % (title_string, date_string, body_string_fragment))
     body_string = body_string[char_count:]
     doc_count += 1

   return doc_count


 def sample_europarl():
   # download europarl.tgz v7, if not already here (in cwd):
   file_name = "europarl.tgz"
   if not os.path.exists(file_name):
     print("Download %s to %s..." % (EUROPARL_V7_URL, file_name))
     urllib.request.urlretrieve(EUROPARL_V7_URL, file_name + ".tmp")
     os.rename(file_name + ".tmp", file_name)
   else:
     print("%s already here; skipping download..." % file_name)

   if not DEBUG:
     tmp_dir_path = tempfile.mkdtemp()
   else:
     tmp_dir_path = "/tmp/tmp31ekzg75"
   print('Using tmp dir "%s"...' % tmp_dir_path)
   try:
     if not DEBUG:
       cmd = "tar xzf %s -C %s" % (file_name, tmp_dir_path)
       print("Run: %s" % cmd)
       subprocess.run(cmd, shell=True, check=False)

     doc_count = 0
     skip_count = 0
     file_count = 0

     all_txt_file_name = "%s/all.txt" % tmp_dir_path

     print("Extract text...")

     start_time = time.time()
     next_print_time = start_time + 3
     # normalize text a bit and concatenate all lines into single file, counting total lines/bytes
     with open(all_txt_file_name, "w", encoding="utf-8") as all_out:
       for dir_path, _, file_names in os.walk("%s/txt" % tmp_dir_path):
         for file_name in file_names:
           if file_name.endswith(".txt"):
             file_count += 1

             year, month, day = (int(x) for x in file_name[3:-4].split("-")[:3])
             if year >= 50:
               year = 1900 + year
             else:
               year = 2000 + year

             date_string = "%04d-%02d-%02d" % (year, month, day)

             # unfortunately we need errors='ignore' since in Europarl v7, one file (pl/ep-09-10-22-009.txt) has invalid utf-8:
             chapter_count = 0
             with open("%s/%s" % (dir_path, file_name), encoding="utf-8", errors="ignore") as f_in:
               last_text: list[str] = []
               last_title = None
               while True:
                 line = f_in.readline()
                 if line == "":
                   break
                 line = line.strip()
                 if line.startswith("<CHAPTER "):
                   if last_title is not None:
                     s = " ".join(last_text)
                     s = re_tag.sub(" ", s)
                     s = re_newlines.sub(" ", s)
                     s = s.strip()
                     if len(s) > 0:
                       doc_count += split_docs(all_out, last_title, date_string, s)
                     else:
                       skip_count += 1

                     last_text = []
                     chapter_count += 1
                   while True:
                     last_title = f_in.readline()
                     if last_title == "":
                       last_title = None
                       break
                     last_title = re_tag.sub(" ", last_title).strip()
                     if len(last_title) > 0:
                       break
                   continue
                 last_text.append(line)

               if last_title is not None:
                 s = " ".join(last_text)
                 s = re_tag.sub(" ", s)
                 s = re_newlines.sub(" ", s)
                 s = s.strip()
                 if len(s) > 0:
                   doc_count += split_docs(all_out, last_title, date_string, s)
                 else:
                   skip_count += 1
                 chapter_count += 1
               else:
                 skip_count += 1

               if chapter_count > 0:
                 # print('%s/%s: %d chapters' % (dir_path, file_name, chapter_count))
                 pass

             now = time.time()
             if now > next_print_time:
               print(
                 "%4.1fs: keep %.2f K of %.2f K files (%.1f%%), %.2f M docs, %.2f GB..."
                 % (now - start_time, (file_count - skip_count) / 1000, file_count / 1000, 100 * (file_count - skip_count) / file_count, doc_count / 1000000, all_out.tell() / 1024 / 1024 / 1024)
               )
               while next_print_time < now:
                 next_print_time += 3

     total_mb = os.path.getsize(all_txt_file_name) / 1024 / 1024
     now = time.time()
     print(
       "%4.1fs (done): keep %.2f K of %.2f K files (%.1f%%), %.2f M docs, %.2f GB..."
       % (
         now - start_time,
         (file_count - skip_count) / 1000,
         file_count / 1000,
         100 * (file_count - skip_count) / file_count,
         doc_count / 1000000,
         os.path.getsize(all_txt_file_name) / 1024 / 1024 / 1024,
       )
     )

     print("Shuffle...")
     subprocess.run("shuf %s > %s.shuffled" % (all_txt_file_name, all_txt_file_name), shell=True, check=False)

     for mb in (20, 200, 2000):
       print("Sample %d MB file..." % mb)
       file_name_out = "%dmb.txt" % mb
       with open(file_name_out, "w", encoding="utf-8") as f_out:
         chance = mb / total_mb

         with open(all_txt_file_name + ".shuffled", encoding="utf-8") as f:
           while True:
             line = f.readline()
             if len(line) == 0:
               break
             if random.random() <= chance:
               f_out.write(line)

       print("  got %.2f MB" % (os.path.getsize(file_name_out) / 1024 / 1024))

       compress_with_seek_points(file_name_out, file_name_out + ".gz", mb)

   finally:
     print('Removing tmp dir "%s"...' % tmp_dir_path)
     if not DEBUG:
       shutil.rmtree(tmp_dir_path)

   print("\nWARNING: left ./europarl.tgz, which you should delete if you do not want it!\n")


 if False:
   compress_with_seek_points("/x/tmp/europarl.lines.txt", "/x/tmp/foo.txt.gz", 16)
 else:
   sample_europarl()
	#!/usr/bin/env python3
	# Licensed to the Apache Software Foundation (ASF) under one or more
	# contributor license agreements. See the NOTICE file distributed with
	# this work for additional information regarding copyright ownership.
	# The ASF licenses this file to You under the Apache License, Version 2.0
	# (the "License"); you may not use this file except in compliance with
	# the License. You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	import gzip
	import os
	import random
	import re
	import shutil
	import subprocess
	import tempfile
	import time
	import urllib.request
	from io import TextIOWrapper

	DEBUG = False

	TARGET_DOC_CHARS = 1024


	def compress_with_seek_points(file_name_in: str, file_name_out: str, num_seek_points: int):
	bytes_per_chunk = os.path.getsize(file_name_in) / num_seek_points

	seek_points: list[int] = []

	if os.path.exists(file_name_out):
	os.remove(file_name_out)

	with open(file_name_in, "rb") as f_in:
	f_out = None

	bytes_in_chunk = 0

	chunk_count = 0

	while True:
	if f_out is None:
	if os.path.exists(file_name_out):
	seek_points.append(os.path.getsize(file_name_out))
	print(" create chunk %s at pos=%s" % (chunk_count, seek_points[-1]))
	else:
	print(" create chunk %s at pos=0" % chunk_count)
	f_out = gzip.open(file_name_out, "ab")
	chunk_count += 1

	line = f_in.readline()
	if len(line) == 0:
	break

	bytes_in_chunk += len(line)
	f_out.write(line)

	if bytes_in_chunk > bytes_per_chunk and chunk_count < num_seek_points:
	f_out.close()
	f_out = None
	bytes_in_chunk = 0

	with open(file_name_out[:-3] + ".seek", "w") as f_out:
	for seek_point in seek_points:
	f_out.write("%d\n" % seek_point)


	re_tag = re.compile(r"<[^>]+?>")
	re_newlines = re.compile(r"\n+")
	re_space = re.compile(r"\s")

	# used to find word break, for splitting docs into ~1 KB sized smaller docs:
	re_next_non_word_character = re.compile(r"\W", re.UNICODE)

	EUROPARL_V7_URL = "https://www.statmt.org/europarl/v7/europarl.tgz"


	def split_docs(all_out: TextIOWrapper, title_string: str, date_string: str, body_string: str):
	"""Splits docs into smallish (~1 KB) sized docs, repeating same title and date"""
	doc_count = 0
	while len(body_string) > 0:
	char_count = int(random.gauss(TARGET_DOC_CHARS, TARGET_DOC_CHARS / 4))
	if char_count < 64:
	# trimmed normal?
	continue

	m = re_next_non_word_character.search(body_string, char_count)
	if m is not None:
	char_count = m.start(0)
	else:
	char_count = len(body_string)

	body_string_fragment = body_string[:char_count].strip()

	# print('write title %d, body %d' % (len(title_string), len(body_string_fragment)))
	all_out.write("%s\t%s\t%s\n" % (title_string, date_string, body_string_fragment))
	body_string = body_string[char_count:]
	doc_count += 1

	return doc_count


	def sample_europarl():
	# download europarl.tgz v7, if not already here (in cwd):
	file_name = "europarl.tgz"
	if not os.path.exists(file_name):
	print("Download %s to %s..." % (EUROPARL_V7_URL, file_name))
	urllib.request.urlretrieve(EUROPARL_V7_URL, file_name + ".tmp")
	os.rename(file_name + ".tmp", file_name)
	else:
	print("%s already here; skipping download..." % file_name)

	if not DEBUG:
	tmp_dir_path = tempfile.mkdtemp()
	else:
	tmp_dir_path = "/tmp/tmp31ekzg75"
	print('Using tmp dir "%s"...' % tmp_dir_path)
	try:
	if not DEBUG:
	cmd = "tar xzf %s -C %s" % (file_name, tmp_dir_path)
	print("Run: %s" % cmd)
	subprocess.run(cmd, shell=True, check=False)

	doc_count = 0
	skip_count = 0
	file_count = 0

	all_txt_file_name = "%s/all.txt" % tmp_dir_path

	print("Extract text...")

	start_time = time.time()
	next_print_time = start_time + 3
	# normalize text a bit and concatenate all lines into single file, counting total lines/bytes
	with open(all_txt_file_name, "w", encoding="utf-8") as all_out:
	for dir_path, _, file_names in os.walk("%s/txt" % tmp_dir_path):
	for file_name in file_names:
	if file_name.endswith(".txt"):
	file_count += 1

	year, month, day = (int(x) for x in file_name[3:-4].split("-")[:3])
	if year >= 50:
	year = 1900 + year
	else:
	year = 2000 + year

	date_string = "%04d-%02d-%02d" % (year, month, day)

	# unfortunately we need errors='ignore' since in Europarl v7, one file (pl/ep-09-10-22-009.txt) has invalid utf-8:
	chapter_count = 0
	with open("%s/%s" % (dir_path, file_name), encoding="utf-8", errors="ignore") as f_in:
	last_text: list[str] = []
	last_title = None
	while True:
	line = f_in.readline()
	if line == "":
	break
	line = line.strip()
	if line.startswith("<CHAPTER "):
	if last_title is not None:
	s = " ".join(last_text)
	s = re_tag.sub(" ", s)
	s = re_newlines.sub(" ", s)
	s = s.strip()
	if len(s) > 0:
	doc_count += split_docs(all_out, last_title, date_string, s)
	else:
	skip_count += 1

	last_text = []
	chapter_count += 1
	while True:
	last_title = f_in.readline()
	if last_title == "":
	last_title = None
	break
	last_title = re_tag.sub(" ", last_title).strip()
	if len(last_title) > 0:
	break
	continue
	last_text.append(line)

	if last_title is not None:
	s = " ".join(last_text)
	s = re_tag.sub(" ", s)
	s = re_newlines.sub(" ", s)
	s = s.strip()
	if len(s) > 0:
	doc_count += split_docs(all_out, last_title, date_string, s)
	else:
	skip_count += 1
	chapter_count += 1
	else:
	skip_count += 1

	if chapter_count > 0:
	# print('%s/%s: %d chapters' % (dir_path, file_name, chapter_count))
	pass

	now = time.time()
	if now > next_print_time:
	print(
	"%4.1fs: keep %.2f K of %.2f K files (%.1f%%), %.2f M docs, %.2f GB..."
	% (now - start_time, (file_count - skip_count) / 1000, file_count / 1000, 100 * (file_count - skip_count) / file_count, doc_count / 1000000, all_out.tell() / 1024 / 1024 / 1024)
	)
	while next_print_time < now:
	next_print_time += 3

	total_mb = os.path.getsize(all_txt_file_name) / 1024 / 1024
	now = time.time()
	print(
	"%4.1fs (done): keep %.2f K of %.2f K files (%.1f%%), %.2f M docs, %.2f GB..."
	% (
	now - start_time,
	(file_count - skip_count) / 1000,
	file_count / 1000,
	100 * (file_count - skip_count) / file_count,
	doc_count / 1000000,
	os.path.getsize(all_txt_file_name) / 1024 / 1024 / 1024,
	)
	)

	print("Shuffle...")
	subprocess.run("shuf %s > %s.shuffled" % (all_txt_file_name, all_txt_file_name), shell=True, check=False)

	for mb in (20, 200, 2000):
	print("Sample %d MB file..." % mb)
	file_name_out = "%dmb.txt" % mb
	with open(file_name_out, "w", encoding="utf-8") as f_out:
	chance = mb / total_mb

	with open(all_txt_file_name + ".shuffled", encoding="utf-8") as f:
	while True:
	line = f.readline()
	if len(line) == 0:
	break
	if random.random() <= chance:
	f_out.write(line)

	print(" got %.2f MB" % (os.path.getsize(file_name_out) / 1024 / 1024))

	compress_with_seek_points(file_name_out, file_name_out + ".gz", mb)

	finally:
	print('Removing tmp dir "%s"...' % tmp_dir_path)
	if not DEBUG:
	shutil.rmtree(tmp_dir_path)

	print("\nWARNING: left ./europarl.tgz, which you should delete if you do not want it!\n")


	if False:
	compress_with_seek_points("/x/tmp/europarl.lines.txt", "/x/tmp/foo.txt.gz", 16)
	else:
	sample_europarl()