Package: tinyllama
~tinyllama~Model classes and pre-training utilities for a tiny version of Llama in PyTorch.Installation 🚀pipinstalltinyllamaPre-training a model 🏋‍♀Initializing a tokenizerWith a simple character-level tokenizer:fromtinyllama.tokenizersimportCharacterTokenizertokenizer=CharacterTokenizer()# '|' is the default eos_tokentokenizer.add_eos_tokens()To turn a corpus into tokens:tokens=tokenizer.tokenize(corpus)Initializing a Llama modelfromtinyllamaimportLlamamodel=Llama(context_window=500,emb_dim=10,n_heads=2,n_blocks=2)Multi-Query attentionmodel=Llama(context_window=500,emb_dim=10,n_heads=2,n_blocks=2,gq_ratio=2)The parameter gq_ratio represents the ratio $\frac{number \ of \ heads}{number \  of \ queries/keys}$, it is set to 1 by default.The configuration above builds a Llama model with the number of heads being twice as much as the number of queries/keys.Launching a pre-training jobfromtinyllamaimportTrainConfig,TrainerTrainConfig=TrainConfig(batch_size=32,epochs=50,log_interval=15)Trainer=Trainer(TrainConfig)Trainer.run(model,tokens)Diagnosis 😷Diagnosis class run a training job on a copy of the model and returns training information that could be useful to the user.Diagnosing the learning rateReturns a plot representing the loss for each learning rate,the scale for the argument start and end is logarithmic.fromtinyllama.diagnosisimportLrDiagnoseLrDiagnose=LrDiagnose(start=-5,end=0,n_lrs=50)# LrDiagnose.run(model, tokens, TrainConfig)LrDiagnose=LrDiagnose(start=-5,end=0,n_lrs=50)LrDiagnose.run(model,tokens,TrainConfig)Diagnosing the gradientsReturns a histogram representing the distribution of the gradients, doesn't run additional training jobs.fromtinyllama.diagnosisimportGradDiagnoseGradDiagnose=GradDiagnose(num_params_to_track=1500)GradDiagnose.run(model)Diagnosing the activation layers (SwiGLU layers)Returns a histogram representing the distribution of the activation layers.fromtinyllama.diagnosisimportSwigluDiagnoseSwigluDiagnose=SwigluDiagnose(num_embeddings_for_histogram=50,track_direction="forward")SwigluDiagnose.run(model,tokens,TrainConfig)Diagnosing the gradients/data ratiosReturns a plot representing the gradient/data ratio in each step of the training.fromtinyllama.diagnosisimportSwigluDiagnoseGdrDiagnose=GdrDiagnose(num_params_to_track=5,num_iters=150)GdrDiagnose.run(model,tokens,TrainConfig)Hyperparameter tuning with GPTune ⚙️GPTune facilitates hyperparameter tuning by leveraging Gaussian Processes as a means to optimize the tuning process.fromtinyllama.gptunerimportGPTuneConfig,GPTuneGPTuneConfig=GPTuneConfig(num_training_samples=100,hyperparams_to_tune=["epochs","n_heads"],l_bounds=[10,2],u_bounds=[50,5],num_evaluations=500)GPTune=GPTune(GPTuneConfig)GPTune.run(model,tokens,TrainConfig)Generating ✍Generates a response to a prompt.fromtinyllamaimportgenerate# kv_cache is set to True by default.generate(model,prompt,max_tokens=900,kv_cache=True)Parsing 📜Parses single or multiple files.# ".txt" filesfromtinyllama.readersimportget_textcorpus=get_text("./txt_path")# ".pdf" filesfromtinyllama.readersimportget_pdf_textcorpus=get_pdf_text("./pdf_path")To parse multiple files:# ".txt" filesfromtinyllama.readersimportget_textcorpus=''.join(get_text(pdf_path)fortxt_pathintxt_paths)# ".pdf" filesfromtinyllama.readersimportget_pdf_textcorpus=''.join(get_pdf_text(pdf_path)forpdf_pathinpdf_paths)