Update app.py

app.py

@@ -3,7 +3,7 @@ import json
 from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments, AutoTokenizer
 from datasets import Dataset
 import shutil
-import os
+import torch
 
 # Load tokenized data
 def load_data(tokenized_file):
@@ -27,8 +27,10 @@ def fine_tune_model(tokenized_file, progress=gr.Progress()):
     training_args = TrainingArguments(
         output_dir='./results',
         num_train_epochs=3,
-        per_device_train_batch_size=16,
-        per_device_eval_batch_size=64,
+        per_device_train_batch_size=8,  # Reduce batch size
+        per_device_eval_batch_size=16,
+        gradient_accumulation_steps=2,  # Use gradient accumulation
+        fp16=True,  # Enable mixed precision training
         warmup_steps=500,
         weight_decay=0.01,
         logging_dir='./logs',