Update app.py

app.py

@@ -18,7 +18,7 @@ import transformers
 import datasets
 from dotenv import load_dotenv
 from datasets import load_dataset, get_dataset_config_names, IterableDataset
-from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, TrainerCallback, AutoConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, TrainerCallback, AutoConfig, DataCollatorForLanguageModeling
 from huggingface_hub import login, whoami, create_repo, upload_folder
 import spaces
 
@@ -190,7 +190,7 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name,
                     if len(batch_buffer) >= 100:
                         encoded_batch = tokenizer(batch_buffer, truncation=True, max_length=2048, padding=False)
                         for input_ids in encoded_batch["input_ids"]:
-                            yield {"input_ids": input_ids, "labels": input_ids}
+                            yield {"input_ids": input_ids}
                         batch_buffer = []
                 except:
                     continue
@@ -212,6 +212,8 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name,
 
         output_dir = f"checkpoints/{job_id}"
         
+        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
         training_args = TrainingArguments(
             output_dir=output_dir,
             per_device_train_batch_size=int(batch_size),
@@ -238,6 +240,7 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name,
             model=original_model,
             train_dataset=dataset_iterable,
             args=training_args,
+            data_collator=data_collator,
             callbacks=[CustomTrainerCallback(job_id, hf_token, full_repo_id)]
         )
 
@@ -343,7 +346,7 @@ def load_from_url(request: gr.Request):
         pass
     return gr.update(selected="launch_tab"), ""
 
-with gr.Blocks(title="Nucleus Enterprise", theme=gr.themes.Base()) as demo:
+with gr.Blocks(title="Nucleus Enterprise") as demo:
     with gr.Column():
         gr.Markdown("# ⚛️ NUCLEUS ENTERPRISE")
         gr.Markdown("Autonomous LLM Foundry | V7.0 Scratch Edition")