Upload app.py

app.py

@@ -6,7 +6,6 @@ from datasets import load_dataset, concatenate_datasets, load_from_disk
 import traceback
 from sklearn.metrics import accuracy_score
 import numpy as np
-import torch
 
 import os
 from huggingface_hub import login
@@ -27,7 +26,6 @@ model = AutoModelForSeq2SeqLM.from_pretrained('google/t5-efficient-tiny', num_la
 @spaces.GPU(duration=120)
 def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
     try:
-        torch.cuda.empty_cache()
         def compute_metrics(eval_pred):
             logits, labels = eval_pred
             predictions = np.argmax(logits, axis=1)
@@ -68,9 +66,9 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             save_total_limit=3, 
         )
         # Check if a checkpoint exists and load it
-        if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
-            print("Loading model from checkpoint...")
-            model = AutoModelForSeq2SeqLM.from_pretrained(training_args.output_dir)        
+       # if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
+          #  print("Loading model from checkpoint...")
+          #  model = AutoModelForSeq2SeqLM.from_pretrained(training_args.output_dir)        
     
         max_length = 128
         try: