def __init__(self, url: str, is_pre_train: bool = False, must_exist: bool = False, is_replica = False):
self.is_pre_train = is_pre_train
if environment.WORLD_RANK == 0 or is_replica:
encoded_path = pathlib.Path(url.replace("sqlite:///", "")).parent
self.length_monitor = monitors.CumulativeHistMonitor(encoded_path, "encoded_kernel_length")
if not self.is_pre_train:
self.token_monitor = monitors.NormalizedFrequencyMonitor(encoded_path, "token_distribution")
self.feature_monitors = {ftype: monitors.CategoricalDistribMonitor(encoded_path, "{}_distribution".format(ftype)) for ftype in extractor.extractors.keys()}
super(EncodedContentFiles, self).__init__(url, Base, must_exist=must_exist)
if environment.WORLD_SIZE > 1 and not is_replica:
# Conduct engine connections to replicated preprocessed chunks.
self.base_path = pathlib.Path(url.replace("sqlite:///", "")).resolve().parent
hash_id = self.base_path.name
try:
tdir = pathlib.Path(FLAGS.local_filesystem).resolve() / hash_id / "node_encoded"
except Exception:
tdir = pathlib.Path("/tmp").resolve() / hash_id / "node_encoded"
distrib.lock()
tdir.mkdir(parents = True, exist_ok = True)
distrib.unlock()
self.replicated_path = tdir / "encoded_{}.db".format(environment.WORLD_RANK)
self.replicated = EncodedContentFiles(
url = "sqlite:///{}".format(str(self.replicated_path)),
is_pre_train = is_pre_train,
must_exist = must_exist,
is_replica = True
)
self.length_monitor = self.replicated.length_monitor
if not self.is_pre_train:
self.token_monitor = self.replicated.token_monitor
self.feature_monitors = self.replicated.feature_monitors
distrib.barrier()
return
def _ConfigTrainParams(self,
data_generator: torchLMDataGenerator,
pre_train: bool,
) -> None:
"""
Model parameter initialization for training and validation.
"""
self._ConfigModelParams(is_sampling = False)
self.train_batch_size = self.config.training.batch_size
self.eval_batch_size = self.config.training.batch_size
self.learning_rate = self.config.training.adam_optimizer.initial_learning_rate_micros / 1e6
self.num_warmup_steps = self.config.training.num_warmup_steps if not pre_train else self.config.training.num_prewarmup_steps
self.max_grad_norm = 1.0
self.steps_per_epoch = data_generator.steps_per_epoch
self.current_step = None
self.num_epochs = data_generator.num_epochs
self.num_train_steps = self.steps_per_epoch * self.num_epochs
self.max_eval_steps = FLAGS.max_eval_steps
self.validation_results_file = "val_results.txt"
self.validation_results_path = os.path.join(str(self.logfile_path if not pre_train else self.pre_logfile_path), self.validation_results_file)
m = model.BertForPreTraining(
self.bert_config,
tokenizer = self.tokenizer,
target_lm = "hole" if self.config.training.data_generator.HasField("hole") else "mask"
).to(self.pytorch.offset_device)
if self.pytorch.num_nodes > 1:
distrib.barrier()
m = self.torch.nn.parallel.DistributedDataParallel(
m,
device_ids = [self.pytorch.offset_device],
output_device = self.pytorch.offset_device,
)
elif self.pytorch.num_gpus > 1:
m = self.torch.nn.DataParallel(m)
opt, lr_scheduler = optimizer.create_optimizer_and_scheduler(
model = m,
num_train_steps = self.num_train_steps,
warmup_steps = self.num_warmup_steps,
learning_rate = self.learning_rate,
)
self.train = torchBert.BertEstimator(
m, data_generator, opt, lr_scheduler
)
l.logger().info(self.GetShortSummary())
return
def tokenizer(self) -> tokenizers.TokenizerBase:
"""Must call Create() first."""
if not self._created:
raise ValueError("Must call Create() before accessing tokenizer property.")
if self._tokenizer is None:
if self.tokenizer_path.is_file():
self._tokenizer = tokenizers.TokenizerBase.FromFile(self.tokenizer_path)
else:
if environment.WORLD_RANK == 0:
self._tokenizer = self._CreateTokenizer()
l.logger().warn("Created tokenizer.")
distrib.barrier()
if environment.WORLD_RANK != 0:
self._tokenizer = tokenizers.TokenizerBase.FromFile(self.tokenizer_path)
return self._tokenizer
def GetContentFileRoot(self, config: corpus_pb2.Corpus) -> pathlib.Path:
"""Get the path of the directory containing content files.
If the corpus is a local directory, this simply returns the path. Otherwise,
this method creates a temporary copy of the files which can be used within
the scope of this context.
Args:
config: The corpus config proto.
Returns:
The path of a directory containing content files.
"""
if config.HasField("local_directory"):
yield pathlib.Path(ExpandConfigPath(config.local_directory))
elif config.HasField("local_tar_archive"):
with tempfile.TemporaryDirectory(prefix="clgen_corpus_", dir = FLAGS.local_filesystem) as d:
l.logger().info("Unpacking {}...".format(ExpandConfigPath(config.local_tar_archive).name))
start_time = time.time()
if environment.WORLD_RANK == 0:
cmd = [
"tar",
"-xf",
str(ExpandConfigPath(config.local_tar_archive)),
"-C",
d,
]
subprocess.check_call(cmd)
distrib.barrier()
l.logger().info(
"Unpacked {} in {} ms".format(
ExpandConfigPath(config.local_tar_archive).name,
humanize.intcomma(int((time.time() - start_time) * 1000)),
)
)
yield pathlib.Path(d)
elif config.HasField("bq_database"):
input_bq = pathlib.Path(ExpandConfigPath(config.bq_database))
if environment.WORLD_SIZE > 1:
target_bq = self.replicated_path.parent / "bq_database_replica_{}.db".format(environment.WORLD_RANK)
if not target_bq.exists():
shutil.copy(input_bq, target_bq)
yield target_bq
else:
yield input_bq
else:
raise NotImplementedError
def MergeReplicas(self) -> None:
"""
When distributed nodes work for the same preprocessed DB
this function moves finalized preprocessed chunks back into the AFS
and master node merges them into the final preprocessed.db
"""
shutil.copy(
self.replicated.url.replace("sqlite:///", ""), self.base_path / "preprocessed_{}.db".format(environment.WORLD_RANK)
)
distrib.barrier()
if environment.WORLD_RANK == 0:
db_chunks = glob.glob(str(self.base_path / "preprocessed_*.db"))
dbs = [PreprocessedContentFiles(url = "sqlite:///{}".format(p), must_exist = True, is_replica = True) for p in db_chunks]
merge_db(dbs, self)
for p in db_chunks:
os.remove(p)
distrib.barrier()
return
def GetHashOfArchiveContents(archive: pathlib.Path) -> str:
"""Compute the checksum of the contents of a directory.
Args:
archive: Path of the archive.
Returns:
Checksum of the archive.
Raises:
UserError: If the requested archive does not exist, or cannot be unpacked.
"""
if not (archive.parent / "corpus_registry.json").exists():
raise FileNotFoundError("corpus_registry.json file not found.")
with open(archive.parent / "corpus_registry.json", 'r') as js:
reg = json.load(js)
if archive.name not in reg:
raise FileNotFoundError("Corpus {} is not registered in corpus_registry".format(archive.name))
if not archive.is_file():
l.logger().info("Corpus found in registry. Downloading from Google Drive...")
if environment.WORLD_RANK == 0:
gdown.download("https://drive.google.com/uc?id={}".format(reg[archive.name]['url']), str(archive))
distrib.barrier()
if 'hash' in reg[archive.name]:
return reg[archive.name]['hash']
else:
with tempfile.TemporaryDirectory(prefix="clgen_corpus_", dir = FLAGS.local_filesystem) as d:
pv = ["pv", str(archive)]
tar = ["tar", "xfj", "-", "-C", d]
try:
pv_proc = subprocess.Popen(pv, stdout = subprocess.PIPE)
subprocess.check_call(tar, stdin = pv_proc.stdout)
except subprocess.CalledProcessError:
raise ValueError(f"Archive unpack failed: '{archive}'")
return checksumdir.dirhash(d, "sha1")
def __init__(self, url: str, must_exist: bool = False, is_replica = False):
if environment.WORLD_RANK == 0 or is_replica:
super(PreprocessedContentFiles, self).__init__(
url, Base, must_exist=must_exist
)
if environment.WORLD_SIZE > 1 and not is_replica:
# Conduct engine connections to replicated preprocessed chunks.
self.base_path = pathlib.Path(url.replace("sqlite:///", "")).resolve().parent
hash_id = self.base_path.name
try:
tdir = pathlib.Path(FLAGS.local_filesystem).resolve() / hash_id / "node_preprocessed"
except Exception:
tdir = pathlib.Path("/tmp").resolve() / hash_id / "node_preprocessed"
distrib.lock()
tdir.mkdir(parents = True, exist_ok = True)
distrib.unlock()
self.replicated_path = tdir / "preprocessed_{}.db".format(environment.WORLD_RANK)
self.replicated = PreprocessedContentFiles(
url = "sqlite:///{}".format(str(self.replicated_path)),
must_exist = must_exist,
is_replica = True
)
distrib.barrier()
return
def createCorpus(self, path: pathlib.Path) -> np.array:
"""
Constructs training corpus in text format, stores it in
shaped_corpus
Each corpus datapoint is either a single kernel or a random
sequence of size sequence_length (legacy).
If corpus has been previously pickled and stored, then it is loaded.
"""
start_time = time.time()
# Set corpus dimension parameters
sequence_length = self.training_opts.sequence_length
effect_seq_length = sequence_length - (2 if self.config.use_start_end else 0)
batch_size = self.training_opts.batch_size
dupe_factor = self.training_opts.dupe_factor
shuffle = self.training_opts.shuffle_corpus_contentfiles_between_epochs
pad = [self.tokenizer.padToken ]
start = [self.tokenizer.startToken ]
end = [self.tokenizer.endToken ]
shaped_corpus = None
corpus_file = "{}corpus.pkl".format("pre_" if self.pre_train else "")
# Monitor counts actual length distribution of kernel instances.
kernel_length_monitor = monitors.FrequencyMonitor(path, "{}kernel_length".format("pre_" if self.pre_train else ""))
# Token frequency distribution monitor.
if not self.pre_train:
feature_monitors = {
ftype: monitors.CategoricalDistribMonitor(
path,
"{}{}_distribution".format("pre_" if self.pre_train else "", ftype)
)
for ftype in extractor.extractors.keys()
}
if (path / corpus_file).exists():
with open(path / corpus_file, 'rb') as infile:
shaped_corpus = pickle.load(infile)
if self.num_train_steps:
self.num_epochs = self.num_train_steps // self.config.steps_per_epoch
self.steps_per_epoch = self.config.steps_per_epoch
l.logger().info(
"Loaded from file corpus of {} examples in {} ms.".format(
humanize.intcomma(len(shaped_corpus)),
humanize.intcomma(int((time.time() - start_time) * 1000)),
)
)
return shaped_corpus
# generate a kernel corpus
if (path / "text_corpus.pkl").exists():
# Only sampler writes a text_corpus.pkl, to do online or active sampling.
# The reason is, corpus is saved in text format, to be picked up with the
# right tokenizer. And that is the model's tokenizer.
with open(path / "text_corpus.pkl", 'rb') as infile:
encoded_corpus = [self.tokenizer.TokenizeString(x) for x in pickle.load(infile)]
else:
if self.pre_train:
if self.num_train_steps:
self.num_epochs = self.num_train_steps // self.config.steps_per_epoch
self.steps_per_epoch = self.config.steps_per_epoch
if environment.WORLD_RANK == 0:
if len(glob.glob(str(path / "pre_corpus_*.pkl"))) > 0:
return []
encoded_corpus = []
cache_lengths = {}
chunk_size = 250000
i, ch_idx = 0, 0
bar = tqdm.tqdm(total = self.corpus.encoded.size, desc = "Chunk pre-train corpus")
pool = multiprocessing.Pool()
l.logger().info("Processing pre-train corpus chunks...")
for dp in pool.imap_unordered(
functools.partial(
_addStartEndPadToken,
tokenizer = pickle.dumps(self.tokenizer),
trunc = effect_seq_length,
seq_len = sequence_length),
self.corpus.GetTrainingDataGenerator()):
if dp:
rlen, enc_kernel = dp
kernel_length_monitor.register(rlen)
encoded_corpus.append(enc_kernel)
i += 1
if i % chunk_size == 0:
encoded_corpus = np.array(encoded_corpus)
corpus_file = "pre_corpus_{}.pkl".format(ch_idx)
cache_lengths[corpus_file] = len(encoded_corpus)
l.logger().info("Storing chunk {}, len: {}".format(ch_idx, encoded_corpus.shape))
with open(path / corpus_file, 'wb') as outf:
pickle.dump(encoded_corpus, outf, protocol = 4)
with open(path / "pre_lengths_cache.json", 'w') as outf:
json.dump(cache_lengths, outf)
ch_idx += 1
encoded_corpus = []
bar.update(1)
if encoded_corpus:
encoded_corpus = np.array(encoded_corpus)
l.logger().info("Storing chunk {}, len: {}".format(ch_idx, encoded_corpus.shape))
corpus_file = "pre_corpus_{}.pkl".format(ch_idx)
cache_lengths[corpus_file] = len(encoded_corpus)
with open(path / corpus_file, 'wb') as outf:
pickle.dump(encoded_corpus, outf, protocol = 4)
with open(path / "pre_lengths_cache.json", 'w') as outf:
json.dump(cache_lengths, outf)
kernel_length_monitor.plot()
pool.close()
distrib.barrier()
else:
distrib.barrier()
if len(glob.glob(str(path / "pre_corpus_*.pkl"))) > 0:
return []
else:
raise FileNotFoundError(glob.glob(str(path / "pre_corpus_*.pkl")))
return encoded_corpus
else:
if environment.WORLD_RANK == 0:
if not self.feature_encoder:
encoded_corpus = self.corpus.GetTrainingData(sequence_length = effect_seq_length if not self.config.truncate_large_kernels else None)
else:
encoded_corpus = self.corpus.GetTrainingDataWFeatures(sequence_length = effect_seq_length if not self.config.truncate_large_kernels else None)
if self.config.datapoint_type == "kernel":
if environment.WORLD_RANK == 0:
# Reject larger than sequence length
initial_length = copy.deepcopy(len(encoded_corpus))
if not self.pre_train:
# Get features of fitting dataset within sequence length
for feature in self.corpus.GetTrainingFeatures(effect_seq_length):
for ftype, fvector in feature.items():
feature_monitors[ftype].register(fvector)
if self.feature_encoder:
training_features = [x for _, x in encoded_corpus]
encoded_corpus = [x for x, _ in encoded_corpus]
idx, t = set(), []
if self.config.truncate_large_kernels:
for i, x in enumerate(encoded_corpus):
if len(x[:effect_seq_length]) <= effect_seq_length:
t.append(list(x[:effect_seq_length]))
else:
idx.add(i)
else:
for i, x in enumerate(encoded_corpus):
if len(x) <= effect_seq_length:
t.append(list(x))
else:
idx.add(i)
encoded_corpus = t
if self.feature_encoder:
training_features = [x for i, x in enumerate(training_features) if i not in idx]
reduced_length = copy.deepcopy(len(encoded_corpus))
# Add start/end tokens
if self.config.use_start_end:
encoded_corpus = [self._addStartEndToken(kf) for kf in encoded_corpus]
# Register the actual lengths before padding.
kernel_length_monitor.register([len(x) for x in encoded_corpus])
# pad sequences to sequence length
encoded_corpus = np.array([x + pad * (sequence_length - len(x)) for x in encoded_corpus])
if self.feature_encoder:
expanded_corpus = []
encoded_features = []
for dp, fvec in zip(encoded_corpus, training_features):
for fspace in extractor.extractors.keys():
if fspace in fvec:
expanded_corpus.append(dp)
encoded_features.append(self.feature_tokenizer.TokenizeFeatureVector(fvec[fspace], fspace, self.feature_sequence_length))
shaped_corpus = [[src, feats] for src, feats in zip(expanded_corpus, encoded_features)]
else:
shaped_corpus = encoded_corpus
# Shuffle
if shuffle:
self.rngen.shuffle(shaped_corpus)
assert len(shaped_corpus) != 0, "Not enought data. All kernels have been rejected."
# Set corpus epoch parameters
if self.num_train_steps:
self.num_epochs = self.num_train_steps // self.config.steps_per_epoch
self.steps_per_epoch = self.config.steps_per_epoch
l.logger().info("{} kernels were rejected (larger than sequence_length)".format(initial_length - reduced_length))
l.logger().info(
"Loaded corpus of shape {}x{} multiplied by dupe factor: {} in {} ms.".format(
len(shaped_corpus),
sequence_length,
dupe_factor,
humanize.intcomma(int((time.time() - start_time) * 1000)),
)
)
else:
# Set corpus epoch parameters
if self.num_train_steps:
self.num_epochs = self.num_train_steps // self.config.steps_per_epoch
self.steps_per_epoch = self.config.steps_per_epoch
elif self.config.datapoint_type == "statement":
## This branch is legacy data processing and does not support DDP.
if shuffle:
self.rngen.shuffle(encoded_corpus)
encoded_corpus = np.concatenate(encoded_corpus)
# encoded_corpus = np.tile(encoded_corpus, dupe_factor)
# Set corpus dimension parameters
self.steps_per_epoch = len(encoded_corpus) // (batch_size * sequence_length * dupe_factor)
assert self.steps_per_epoch != 0, "Not enought data. Use smaller sequence_length and/or batch_size"
if self.num_train_steps:
self.num_epochs = self.num_train_steps // self.steps_per_epoch
# clipped_corpus_length = dupe_factor * self.steps_per_epoch * batch_size * sequence_length
clipped_corpus_length = self.steps_per_epoch * batch_size * sequence_length
clipped_corpus = encoded_corpus[:clipped_corpus_length]
# shaped_corpus = np.split(clipped_corpus, batch_size * self.steps_per_epoch * dupe_factor, 0)
shaped_corpus = np.split(clipped_corpus, batch_size * self.steps_per_epoch, 0)
# Register the actual lengths before padding.
kernel_length_monitor.register([len(x) for x in shaped_corpus])
np_corpus = np.asarray(shaped_corpus)
assert np_corpus.ndim == 2, "Wrong dimensions for shaped_corpus: {}".format(np_corpus.shape)
assert np_corpus.shape[1] == sequence_length, "Second dimension is not equal to sequence length: {}".format(np_corpus.shape[1])
l.logger().info(
"Loaded corpus of {} tokens (clipped last {} tokens) in {} ms.".format(
humanize.intcomma(clipped_corpus_length),
humanize.intcomma(len(encoded_corpus) - clipped_corpus_length),
humanize.intcomma(int((time.time() - start_time) * 1000)),
)
)
else:
raise ValueError("Unrecognized datapoint_type: {}".format(self.config.datapoint_type))
if environment.WORLD_RANK == 0:
kernel_length_monitor.plot()
if not self.pre_train:
for fm in feature_monitors.values():
fm.plot()
with open(path / corpus_file, 'wb') as outf:
pickle.dump(shaped_corpus, outf)
distrib.barrier()
return shaped_corpus
def __init__(self, config: typing.Union[corpus_pb2.Corpus, corpus_pb2.PreTrainCorpus]):
"""Instantiate a corpus from a proto config.
If this is a new corpus, a number of files will be created, which may
take some time.
Args:
config: A Corpus message.
Raises:
TypeError: If the config argument is not a Sampler proto.
UserError: In case the corpus is not found, or config contains invalid
options.
EmptyCorpusException: In case the corpus contains no data.
"""
if not isinstance(config, corpus_pb2.Corpus) and not isinstance(config, corpus_pb2.PreTrainCorpus):
raise TypeError(f"Config must be a Corpus proto. Received: '{type(config).__name__}'")
# Make a local copy of the configuration.
if isinstance(config, corpus_pb2.Corpus):
self.config = corpus_pb2.Corpus()
self.pre_train = False
else:
self.config = corpus_pb2.PreTrainCorpus()
self.pre_train = True
self.config.CopyFrom(AssertConfigIsValid(config))
self._tokenizer = None
self._created = False
# An in-memory cache of the encoded contentfiles indices arrays.
# Set and used in GetTrainingData().
self._indices_arrays: typing.Optional[typing.List[np.array]] = None
if environment.WORLD_RANK == 0:
cache.cachepath("corpus").mkdir(parents=True, exist_ok=True)
distrib.barrier()
self.content_id = ResolveContentId(self.config)
# Database of pre-processed files.
preprocessed_id = ResolvePreprocessedId(self.content_id, self.config)
if environment.WORLD_RANK == 0:
cache.cachepath("corpus", "preprocessed", preprocessed_id).mkdir(exist_ok=True, parents=True)
distrib.barrier()
preprocessed_db_path = cache.cachepath("corpus", "preprocessed",
preprocessed_id, "preprocessed.db")
if self.config.HasField("content_id") and not preprocessed_db_path.is_file():
raise ValueError(f"Content ID not found: '{self.content_id}'")
self.preprocessed = preprocessed.PreprocessedContentFiles(
f"sqlite:///{preprocessed_db_path}"
)
# Create symlink to contentfiles.
if environment.WORLD_RANK == 0:
symlink = (pathlib.Path(self.preprocessed.url[len("sqlite:///") :]).parent / "contentfiles")
if not symlink.is_symlink():
if config.HasField("local_directory"):
os.symlink(
str(ExpandConfigPath(config.local_directory, path_prefix=FLAGS.clgen_local_path_prefix)),
symlink,
)
elif config.HasField("local_tar_archive"):
os.symlink(
str(ExpandConfigPath(config.local_tar_archive, path_prefix=FLAGS.clgen_local_path_prefix)),
symlink,
)
elif config.HasField("bq_database"):
os.symlink(
str(ExpandConfigPath(config.bq_database, path_prefix=FLAGS.clgen_local_path_prefix)),
symlink,
)
# elif config.HasField("fetch_github"):
# os.symlink(
# str(ExpandConfigPath(config.fetch_github, path_prefix=FLAGS.clgen_local_path_prefix)),
# symlink,
# )
distrib.barrier()
# Data of encoded pre-preprocessed files.
encoded_id = ResolveEncodedId(self.content_id, self.config)
if environment.WORLD_RANK == 0:
cache.cachepath("corpus", "encoded", encoded_id).mkdir(exist_ok=True, parents=True)
distrib.barrier()
db_path = cache.cachepath("corpus", "encoded", encoded_id, "encoded.db")
if self.config.HasField("pre_encoded_corpus_url"):
self.encoded = encoded.EncodedContentFiles(config.pre_encoded_corpus_url, self.pre_train)
else:
self.encoded = encoded.EncodedContentFiles(f"sqlite:///{db_path}", self.pre_train)
self.tokenizer_path = cache.cachepath(
"corpus", "encoded", encoded_id, "tokenizer.pkl"
)
if environment.WORLD_RANK == 0 and not self.config.HasField("pre_encoded_corpus_url"):
symlink = (pathlib.Path(self.encoded.url[len("sqlite:///") :]).parent / "preprocessed")
if not symlink.is_symlink():
os.symlink(
os.path.relpath(
pathlib.Path(self.preprocessed.url[len("sqlite:///") :]).parent,
pathlib.Path(self.encoded.url[len("sqlite:///") :]).parent,
),
symlink,
)
self.hash = encoded_id
self.cache = cache.mkcache("corpus", "encoded", encoded_id)
if environment.WORLD_RANK == 0:
commit.saveCommit(self.cache.path)
commit.saveCommit(self.cache.path.parent.parent / "preprocessed" / preprocessed_id)
distrib.barrier()
l.logger().info("Initialized {}train corpus in {}".format("pre_" if self.pre_train else "", self.cache.path))
return