Update app.py

app.py

@@ -30,7 +30,6 @@ from datasets import Dataset
 from transformers import AutoTokenizer
 import unsloth
 
-unsloth.optimize_all()
 
 model_name = "meta-llama/Llama-2-7b-hf"  # You can use a smaller one like "meta-llama/Llama-2-7b-chat-hf"
 
@@ -46,6 +45,7 @@ tokenizer.pad_token = tokenizer.eos_token  # Use EOS as PAD token
 # tokenizer.add_special_tokens({'pad_token': '[PAD]'})
 
 # Tokenization function
+@unsloth.optimize
 def tokenize_function(examples):
     tokens = tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)
     tokens["labels"] = tokens["input_ids"].copy()  # Use input as labels for text generation
@@ -95,6 +95,7 @@ trainer = Trainer(
     tokenizer=tokenizer,
 )
 
+@unsloth.optimize
 trainer.train()
 
 model.save_pretrained("./fine_tuned_llama2")