Update app.py

app.py

@@ -1,5 +1,3 @@
-print('Importing goodies..')
-
 import gc
 
 import numpy as np
@@ -12,12 +10,8 @@ from datasets import load_dataset
 from tokenizers import ByteLevelBPETokenizer
 import trl
 
-print("Loading dataset..")
-
 dataset = load_dataset("nroggendorff/openhermes", split="train").select(range(int(1e+4)))
 
-print("Setting up tokenizer..")
-
 def get_training_corpus():
     for i in range(0, len(dataset), 1000):
         yield dataset[i : i + 1000]["text"]
@@ -64,8 +58,6 @@ tokenizer.save_pretrained("/tmp/llama-tokenizer")
 tokenizer = AutoTokenizer.from_pretrained("/tmp/llama-tokenizer")
 print(tokenizer.apply_chat_template([{"role": "user", "content": "Why is the sky blue?"}, {"role": "assistant", "content": "Due to rayleigh scattering."}, {"role": "user", "content": "That's cool."}, {"role": "assistant", "content": "Yeah, I agree."}], tokenize=False))
 
-print("Configuring..")
-
 config = LlamaConfig(
     vocab_size=tokenizer.vocab_size,
     hidden_size=int(512 / 1),
@@ -84,8 +76,6 @@ config = LlamaConfig(
 
 model = LlamaForCausalLM(config)
 
-print("Mapping dataset..")
-
 def format_prompts(examples):
     texts = []
     for text in examples['text']:
@@ -106,8 +96,6 @@ dataset = dataset.map(format_prompts, batched=True)
 
 print(dataset['text'][2])
 
-print("Defining trainer..")
-
 args = TrainingArguments(
     output_dir="mayo",
     num_train_epochs=1,
@@ -134,12 +122,8 @@ torch.cuda.set_device(0)
 gc.collect()
 torch.cuda.empty_cache()
 
-print("Training..")
-
 trainer.train()
 
-print("Pushing to hub..")
-
 #trainer.push_to_hub()
 trained_model = trainer.model
 trained_tokenizer = trainer.tokenizer