resolving bugs

app.py

@@ -107,7 +107,7 @@ class HTMLDataset(Dataset):
         }
 
 
-def train_model(progress=gr.Progress()):
+def train_model():
     # Generate synthetic dataset
     dataset = generate_dataset(num_samples=1000)
     train_data, val_data = train_test_split(dataset, test_size=0.2, random_state=42)
@@ -129,10 +129,10 @@ def train_model(progress=gr.Progress()):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model.to(device)
 
-    for epoch in progress.tqdm(range(EPOCHS), desc="Training Progress"):
+    for epoch in range(EPOCHS):
         model.train()
         train_loss = 0
-        for batch in train_dataloader:
+        for batch in tqdm(train_dataloader, desc=f"Epoch {epoch + 1}/{EPOCHS}"):
             input_ids = batch['input_ids'].to(device)
             attention_mask = batch['attention_mask'].to(device)
             labels = batch['labels'].to(device)
@@ -160,7 +160,7 @@ def train_model(progress=gr.Progress()):
         avg_train_loss = train_loss / len(train_dataloader)
         avg_val_loss = val_loss / len(val_dataloader)
 
-        progress(f"Epoch {epoch + 1}/{EPOCHS}, Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}")
+        print(f"Epoch {epoch + 1}/{EPOCHS}, Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}")
 
     return model, tokenizer