SYMBOL INDEX (283 symbols across 28 files)

FILE: arguments.py
  function add_model_args (line 22) | def add_model_args(parser: argparse.ArgumentParser):
  function add_runtime_args (line 43) | def add_runtime_args(parser: argparse.ArgumentParser):
  function add_data_args (line 68) | def add_data_args(parser: argparse.ArgumentParser):
  function add_hp_args (line 99) | def add_hp_args(parser: argparse.ArgumentParser):
  function add_ppo_args (line 148) | def add_ppo_args(parser: argparse.ArgumentParser):
  function add_minillm_args (line 163) | def add_minillm_args(parser: argparse.ArgumentParser):
  function add_distillm_args (line 174) | def add_distillm_args(parser: argparse.ArgumentParser):
  function add_gen_args (line 197) | def add_gen_args(parser: argparse.ArgumentParser):
  function add_peft_args (line 211) | def add_peft_args(parser: argparse.ArgumentParser):
  function get_args (line 225) | def get_args():

FILE: data_utils/distributed_indexed.py
  function code (line 41) | def code(dtype):
  function index_file_path (line 48) | def index_file_path(prefix_path):
  function data_file_path (line 52) | def data_file_path(prefix_path):
  class DistributedMMapIndexedDataset (line 56) | class DistributedMMapIndexedDataset(torch.utils.data.Dataset):
    class Index (line 57) | class Index(object):
      method __init__ (line 59) | def __init__(self, path):
      method __del__ (line 89) | def __del__(self):
      method dtype (line 94) | def dtype(self):
      method sizes (line 98) | def sizes(self):
      method doc_idx (line 102) | def doc_idx(self):
      method __getitem__ (line 105) | def __getitem__(self, i):
      method __len__ (line 108) | def __len__(self):
    method __init__ (line 111) | def __init__(self, path, name, rank_number, rank_total, cache = None):
    method _probe_data_path (line 133) | def _probe_data_path(self, path, name, rank_total):
    method __getstate__ (line 150) | def __getstate__(self):
    method __setstate__ (line 153) | def __setstate__(self, state):
    method _do_init (line 157) | def _do_init(self, path, name, cache, state):
    method __del__ (line 171) | def __del__(self):
    method __len__ (line 178) | def __len__(self):
    method _next_file (line 181) | def _next_file(self):
    method __relative_idx (line 188) | def __relative_idx(self, idx):
    method __slice_item (line 192) | def __slice_item(self, start, stop):
    method __getitem__ (line 199) | def __getitem__(self, idx):
    method sizes (line 209) | def sizes(self):
    method exists (line 212) | def exists(self, path):

FILE: data_utils/indexed_dataset.py
  function __best_fitting_dtype (line 23) | def __best_fitting_dtype(vocab_size=None):
  function get_available_dataset_impl (line 30) | def get_available_dataset_impl():
  function infer_dataset_impl (line 34) | def infer_dataset_impl(path):
  function make_builder (line 50) | def make_builder(out_file, impl, dtype):
  function make_dataset (line 57) | def make_dataset(path, impl, skip_warmup=False):
  function dataset_exists (line 74) | def dataset_exists(path, impl):
  function read_longs (line 81) | def read_longs(f, n):
  function write_longs (line 87) | def write_longs(f, a):
  function code (line 104) | def code(dtype):
  function index_file_path (line 111) | def index_file_path(prefix_path):
  function data_file_path (line 115) | def data_file_path(prefix_path):
  function create_doc_idx (line 119) | def create_doc_idx(sizes):
  class IndexedDataset (line 127) | class IndexedDataset(torch.utils.data.Dataset):
    method __init__ (line 131) | def __init__(self, path):
    method read_index (line 137) | def read_index(self, path):
    method read_data (line 155) | def read_data(self, path):
    method check_index (line 158) | def check_index(self, i):
    method __del__ (line 162) | def __del__(self):
    method __getitem__ (line 167) | def __getitem__(self, idx):
    method __len__ (line 191) | def __len__(self):
    method num_tokens (line 194) | def num_tokens(self, index):
    method size (line 197) | def size(self, index):
    method exists (line 201) | def exists(path):
    method supports_prefetch (line 207) | def supports_prefetch(self):
  class IndexedCachedDataset (line 211) | class IndexedCachedDataset(IndexedDataset):
    method __init__ (line 213) | def __init__(self, path):
    method supports_prefetch (line 219) | def supports_prefetch(self):
    method prefetch (line 222) | def prefetch(self, indices):
    method __getitem__ (line 247) | def __getitem__(self, idx):
  class IndexedDatasetBuilder (line 264) | class IndexedDatasetBuilder(object):
    method __init__ (line 275) | def __init__(self, out_file, dtype=np.int32):
    method add_item (line 284) | def add_item(self, tensor):
    method end_document (line 291) | def end_document(self):
    method merge_file_ (line 294) | def merge_file_(self, another_file):
    method finalize (line 314) | def finalize(self, index_file):
  function _warmup_mmap_file (line 329) | def _warmup_mmap_file(path):
  class MMapIndexedDataset (line 335) | class MMapIndexedDataset(torch.utils.data.Dataset):
    class Index (line 336) | class Index(object):
      method writer (line 340) | def writer(cls, path, dtype):
      method __init__ (line 385) | def __init__(self, path, skip_warmup=False):
      method __del__ (line 422) | def __del__(self):
      method dtype (line 427) | def dtype(self):
      method sizes (line 431) | def sizes(self):
      method doc_idx (line 435) | def doc_idx(self):
      method __getitem__ (line 439) | def __getitem__(self, i):
      method __len__ (line 442) | def __len__(self):
    method __init__ (line 445) | def __init__(self, path, skip_warmup=False):
    method __getstate__ (line 454) | def __getstate__(self):
    method __setstate__ (line 457) | def __setstate__(self, state):
    method _do_init (line 460) | def _do_init(self, path, skip_warmup):
    method __del__ (line 472) | def __del__(self):
    method __len__ (line 477) | def __len__(self):
    method __getitem__ (line 481) | def __getitem__(self, idx):
    method get (line 501) | def get(self, idx, offset=0, length=None):
    method sizes (line 516) | def sizes(self):
    method supports_prefetch (line 530) | def supports_prefetch(self):
    method exists (line 534) | def exists(path):
  class MMapIndexedDatasetBuilder (line 540) | class MMapIndexedDatasetBuilder(object):
    method __init__ (line 541) | def __init__(self, out_file, dtype=np.int64):
    method add_item (line 547) | def add_item(self, tensor):
    method end_document (line 552) | def end_document(self):
    method merge_file_ (line 555) | def merge_file_(self, another_file):
    method finalize (line 567) | def finalize(self, index_file):

FILE: data_utils/lm_datasets.py
  class LMTrainDataset (line 15) | class LMTrainDataset(Dataset):
    method __init__ (line 16) | def __init__(self, args, tokenizer, path, split, num, ratio, rng_sampl...
    method __len__ (line 40) | def __len__(self):
    method __getitem__ (line 43) | def __getitem__(self, index):
    method _get_lm (line 46) | def _get_lm(self, index):
    method _process_lm (line 53) | def _process_lm(self, i, samp, model_data, no_model_data, gen_data):
    method move_to_device (line 77) | def move_to_device(self, model_data, no_model_data, gen_data, device):
    method collate (line 89) | def collate(self, samples):

FILE: data_utils/prompt_datasets.py
  class PromptDataset (line 13) | class PromptDataset(Dataset):
    method __init__ (line 14) | def __init__(self, args, tokenizer, split, data_path=None, num=-1):
    method __len__ (line 50) | def __len__(self):
    method load_data_json (line 53) | def load_data_json(self, data_path):
    method load_data_txt (line 80) | def load_data_txt(self, data_path):
    method verbalizer (line 93) | def verbalizer(self):
    method __getitem__ (line 96) | def __getitem__(self, index: int):
    method collate (line 114) | def collate(self, samples):
    method move_to_device (line 141) | def move_to_device(self, model_batch, no_model_batch, device):

FILE: distillm/buffer.py
  class ReplayBuffer (line 16) | class ReplayBuffer:
    method __init__ (line 17) | def __init__(self, args):
    method __len__ (line 28) | def __len__(self):
    method sample (line 31) | def sample(self):
    method move_to_device (line 54) | def move_to_device(self, model_data, no_model_data, device):
    method move_to_memory (line 63) | def move_to_memory(self, model_data, no_model_data):

FILE: distillm/losses.py
  function forward_kl (line 4) | def forward_kl(logits, teacher_logits, no_model_batch):
  function reverse_kl (line 14) | def reverse_kl(logits, teacher_logits, no_model_batch):
  function symmetric_kl (line 26) | def symmetric_kl(logits, teacher_logits, no_model_batch, lam=0.9):
  function js_distance (line 32) | def js_distance(logits, teacher_logits, no_model_batch, lam=0.9):
  function tv_distance (line 55) | def tv_distance(logits, teacher_logits, no_model_batch):
  function skewed_forward_kl (line 66) | def skewed_forward_kl(logits, teacher_logits, no_model_batch, lam=0.1):
  function skewed_reverse_kl (line 80) | def skewed_reverse_kl(logits, teacher_logits, no_model_batch, lam=0.1):

FILE: distillm/sampler.py
  class SampleGenerator (line 6) | class SampleGenerator():
    method __init__ (line 7) | def __init__(self, args, tokenizer):
    method run_sample (line 26) | def run_sample(self, model, gen_data):

FILE: evaluate.py
  function setup_model (line 23) | def setup_model(args, ds_config, device):
  function main (line 44) | def main():

FILE: evaluate_main.py
  function prepare_dataset_main (line 22) | def prepare_dataset_main(args, tokenizer):
  function run_model (line 29) | def run_model(args, tokenizer, model, dataset: PromptDataset, epoch, dev...
  function evaluate_main (line 124) | def evaluate_main(args, tokenizer, model, dataset: PromptDataset, split,...

FILE: finetune.py
  function get_teacher_model (line 50) | def get_teacher_model(args, device):
  function get_optimizer (line 77) | def get_optimizer(args, model):
  function get_learning_rate_scheduler (line 95) | def get_learning_rate_scheduler(args, optimizer):
  function setup_model_and_optimizer (line 119) | def setup_model_and_optimizer(args, ds_config, device, set_optim=True):
  function prepare_dataset (line 143) | def prepare_dataset(args, tokenizer):
  function pt_loss (line 162) | def pt_loss(args, model, model_batch, no_model_batch):
  function get_distil_loss (line 171) | def get_distil_loss(args, tokenizer, model, teacher_model, model_batch, ...
  function get_teacher_lm_loss (line 196) | def get_teacher_lm_loss(args, tokenizer, model, teacher_model, model_bat...
  function finetune (line 238) | def finetune(args, tokenizer: AutoTokenizer, model: deepspeed.DeepSpeedE...
  function evaluate (line 431) | def evaluate(args, tokenizer, model, dataset: LMTrainDataset, split, epo...
  function main (line 538) | def main():

FILE: generate.py
  function setup_model (line 28) | def setup_model(args, ds_config, device):
  function prepare_dataset (line 48) | def prepare_dataset(args, tokenizer):
  function generate (line 55) | def generate(args, tokenizer, model, dataset, device):
  function main (line 123) | def main():

FILE: minillm/__init__.py
  function train (line 10) | def train(

FILE: minillm/data_types.py
  class PromptElement (line 7) | class PromptElement:
  class PromptBatch (line 23) | class PromptBatch:
  class PPORLElement (line 39) | class PPORLElement:
  class PPORLBatch (line 80) | class PPORLBatch:

FILE: minillm/losses.py
  class Loss (line 15) | class Loss():
    method __init__ (line 16) | def __init__(self, args, trainer):
    method _get_cumsum_rewards (line 20) | def _get_cumsum_rewards(self, rewards):
    method _get_advantages_and_returns (line 27) | def _get_advantages_and_returns(
    method _pg_loss (line 58) | def _pg_loss(
    method _reg_loss (line 98) | def _reg_loss(self, query_ids, response_ids, mask, logits, inf_mask, s...
    method get_input_batch (line 110) | def get_input_batch(self, ppo_batch: PPORLBatch, pt_batch):
    method ppo_loss (line 122) | def ppo_loss(self, batch: PPORLBatch, logits):
    method pt_loss (line 194) | def pt_loss(self, batch, logits):

FILE: minillm/model.py
  class PPOModel (line 8) | class PPOModel(nn.Module):
    method __init__ (line 9) | def __init__(self, args, device):
    method forward (line 16) | def forward(self, **x):
    method generate (line 20) | def generate(self, **x):
    method set_force_gradient_checkpointing (line 23) | def set_force_gradient_checkpointing(self, value):

FILE: minillm/pipelines.py
  class PPOPipeline (line 15) | class PPOPipeline():
    method __init__ (line 16) | def __init__(self, args, tokenizer, split, ppo_data_path=None, fix_pro...
    method __len__ (line 41) | def __len__(self):
    method __getitem__ (line 44) | def __getitem__(self, index: int):
    method collate (line 60) | def collate(self, samples):
    method move_to_device (line 88) | def move_to_device(self, model_batch, no_model_batch, device):
    method create_loader (line 96) | def create_loader(self, batch_size: int, shuffle=False, drop_last: boo...
  class LMPipeline (line 110) | class LMPipeline():
    method __init__ (line 111) | def __init__(self, args, tokenizer, split, lm_data_path=None, num=-1):
    method __len__ (line 126) | def __len__(self):
    method __getitem__ (line 129) | def __getitem__(self, index):
    method _get_lm (line 132) | def _get_lm(self, index):
    method _process_lm (line 139) | def _process_lm(self, i, samp, model_data, no_model_data):
    method move_to_device (line 157) | def move_to_device(self, model_batch, no_model_batch, device):
    method collate (line 166) | def collate(self, samples):
    method create_loader (line 189) | def create_loader(self, batch_size: int, shuffle=False, drop_last: boo...

FILE: minillm/reward.py
  class Reward (line 8) | class Reward():
    method __init__ (line 9) | def __init__(self, args, tokenizer: AutoTokenizer, model: AutoModelFor...
    method get_input_batch (line 16) | def get_input_batch(self, input_ids, gen_ids, output_pos=True):
    method reward_fn (line 33) | def reward_fn(self, input_ids, gen_ids, inf_mask=None, output_pos=True):

FILE: minillm/sampler.py
  class PPOSampler (line 12) | class PPOSampler():
    method __init__ (line 18) | def __init__(
    method run_sample (line 39) | def run_sample(self, num_rollouts_per_device: int = 1024, iter_count: ...

FILE: minillm/storages.py
  class BaseRolloutStore (line 17) | class BaseRolloutStore(Dataset):
    method __init__ (line 18) | def __init__(self, capacity=-1):
    method push (line 23) | def push(self, exps: Iterable[Any]):
    method __getitem__ (line 29) | def __getitem__(self, index: int) -> PPORLElement:
    method __len__ (line 32) | def __len__(self) -> int:
    method create_loader (line 36) | def create_loader(
    method broadcast (line 53) | def broadcast(self, batch, src=0, group=None):
    method move_to_device (line 57) | def move_to_device(self, batch, device):
  class PPORolloutStorage (line 61) | class PPORolloutStorage(BaseRolloutStore):
    method __init__ (line 66) | def __init__(self, pad_token_id, seed):
    method push (line 74) | def push(self, exps: Iterable[PPORLElement]):
    method save (line 77) | def save(self, path):
    method load (line 85) | def load(self, path):
    method clear_history (line 89) | def clear_history(self):
    method export_history (line 92) | def export_history(self, location: str):
    method __getitem__ (line 104) | def __getitem__(self, index: int) -> PPORLElement:
    method __len__ (line 107) | def __len__(self) -> int:
    method collate (line 110) | def collate(self, elems: Iterable[PPORLElement]):
    method create_loader (line 170) | def create_loader(self, batch_size: int, shuffle=False, drop_last: boo...
    method broadcast (line 177) | def broadcast(self, batch: PPORLBatch, src=0, group=None):
    method move_to_device (line 181) | def move_to_device(self, batch: PPORLBatch, device):

FILE: minillm/trainer.py
  class PPOTrainer (line 43) | class PPOTrainer():
    method __init__ (line 48) | def __init__(self, args, tokenizer: AutoTokenizer, reward_fn, ds_config):
    method set_teacher_model (line 94) | def set_teacher_model(self, model):
    method set_sampler (line 97) | def set_sampler(self, sampler):
    method setup_optimizer (line 100) | def setup_optimizer(self):
    method setup_scheduler (line 114) | def setup_scheduler(self):
    method setup_ds (line 128) | def setup_ds(self, model, optimizer=None, scheduler=None):
    method add_eval_pipeline (line 139) | def add_eval_pipeline(self, eval_pipeline: PPOPipeline):
    method add_lm_pipeline (line 143) | def add_lm_pipeline(self, lm_pipeline: LMPipeline, eval_lm_pipeline: L...
    method get_model_inputs (line 147) | def get_model_inputs(
    method get_mask (line 170) | def get_mask(self, tokens):
    method forward_model (line 176) | def forward_model(self, batch):
    method compute_logits_and_log_probs (line 184) | def compute_logits_and_log_probs(self, query_ids, response_ids, inf_ma...
    method train (line 220) | def train(self):
    method post_backward_callback (line 372) | def post_backward_callback(self):
    method post_epoch_callback (line 375) | def post_epoch_callback(self, epoch):
    method prepare_learning (line 382) | def prepare_learning(self):
    method evaluate (line 406) | def evaluate(self):
    method evaluate_ppo (line 426) | def evaluate_ppo(self):  # noqa: C901
    method evaluate_pt (line 522) | def evaluate_pt(self):
    method save (line 550) | def save(self, directory: Optional[str] = None):
    method save_evals (line 565) | def save_evals(self, preds, results, response_texts, directory: Option...
    method push_to_store (line 579) | def push_to_store(self, data):
    method generate (line 582) | def generate(self, input_ids, attention_mask=None, mode="base", teache...

FILE: minillm/utils.py
  function get_entropy (line 19) | def get_entropy(gen_logits, inf_mask, mask, model_parallel=False):
  function get_log_probs (line 32) | def get_log_probs(logits, ids, mask, inf_mask=None, model_parallel=False):
  function get_x_entropy (line 48) | def get_x_entropy(logits_1, logits_2, inf_mask, mask, model_parallel=Fal...
  function get_rev_kl (line 61) | def get_rev_kl(log_p, log_q, mask):
  function get_global_statistics (line 67) | def get_global_statistics(xs: torch.Tensor) -> Tuple[float, float, int]:
  function whiten (line 82) | def whiten(xs: torch.Tensor, shift_mean=True, distributed=True) -> torch...
  function significant (line 95) | def significant(x: Number, ndigits=2) -> Number:
  class OptimizerName (line 108) | class OptimizerName(str, Enum):
  function get_optimizer_class (line 118) | def get_optimizer_class(name: OptimizerName):
  class SchedulerName (line 138) | class SchedulerName(str, Enum):
  function get_scheduler_class (line 145) | def get_scheduler_class(name: SchedulerName):

FILE: rouge_metric.py
  function normalize_answer (line 12) | def normalize_answer(s):
  function exact_match (line 28) | def exact_match(prediction, ground_truth, xlingual=False):
  function rouge (line 32) | def rouge(prediction, ground_truth, xlingual=False):
  function metric_max_over_ground_truths (line 38) | def metric_max_over_ground_truths(metric_fn, prediction, ground_truths, ...
  function compute_metrics (line 46) | def compute_metrics(predictions, references, xlingual=False):
  function compute_grouped_metrics (line 69) | def compute_grouped_metrics(predictions, references, groups, xlingual=Fa...
  function parse_args (line 87) | def parse_args():

FILE: tools/convert_mp.py
  function main (line 24) | def main():

FILE: tools/process_data_dolly.py
  class Encoder (line 15) | class Encoder(object):
    method __init__ (line 16) | def __init__(self, args):
    method initializer (line 19) | def initializer(self):
    method encode (line 22) | def encode(self, line):
  function main (line 65) | def main():

FILE: tools/process_data_pretrain.py
  class Encoder (line 14) | class Encoder(object):
    method __init__ (line 15) | def __init__(self, args):
    method initializer (line 18) | def initializer(self):
    method encode (line 21) | def encode(self, line):
  function main (line 28) | def main():

FILE: train_minillm.py
  function get_teacher_model (line 19) | def get_teacher_model(args, device):
  function main (line 44) | def main():

FILE: utils.py
  function print_args (line 24) | def print_args(args):
  function save_rank (line 33) | def save_rank(log_str, save_path, rank=0):
  function print_rank (line 39) | def print_rank(*args, rank=0, **kwargs):
  function all_gather (line 45) | def all_gather(t, dim=0, world_size=None, group=None, op="cat"):
  function set_random_seed (line 58) | def set_random_seed(seed, mp=False):
  function init_distributed (line 69) | def init_distributed(args):
  function init_distributed_ds (line 86) | def init_distributed_ds(args):
  function initialize (line 103) | def initialize(args):
  function get_model (line 120) | def get_model(args, device):
  function get_optimizer_params (line 176) | def get_optimizer_params(args, model: nn.Module):
  function get_optimizer_params_peft (line 190) | def get_optimizer_params_peft(args, model: nn.Module):
  function get_tokenizer (line 200) | def get_tokenizer(args):
  function load_parallel (line 212) | def load_parallel(model, load_dir):
  function save_parallel (line 222) | def save_parallel(model, save_dir):