Update app.py

app.py

@@ -1,91 +1,93 @@
-import time
-import torch
-from transformers import BertForSequenceClassification, AdamW
-from torch.utils.data import DataLoader, TensorDataset
-from transformers import BertTokenizer
-import gradio as gr
-import pandas as pd
-import os
-import spaces
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(device)
-
-model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
-tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-model.to(device)
-
-optimizer = AdamW(model.parameters(), lr=1e-5)
-
-global_data = None
-
-def load_data(file):
-    global global_data
-    df = pd.read_csv(file)
-    inputs = tokenizer(df['text'].tolist(), padding=True, truncation=True, return_tensors="pt")  # Mã hóa văn bản
-    labels = torch.tensor(df['lable'].tolist()).long()  #
-    global_data = TensorDataset(inputs['input_ids'], inputs['attention_mask'], labels)
-
-    print(global_data)
-
-def get_dataloader(start, end, batch_size=8):
-    global global_data
-    subset = torch.utils.data.Subset(global_data, range(start, end))
-    return DataLoader(subset, batch_size=batch_size)
-
-@spaces.GPU(duration=120)
-def train_batch(dataloader):
-    model.train()
-    start_time = time.time()
-
-    for step, batch in enumerate(dataloader):
-        input_ids, attention_mask, labels = batch
-        input_ids, attention_mask, labels = input_ids.to(device), attention_mask.to(device), labels.to(device)
-
-        optimizer.zero_grad()
-        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
-        loss = outputs.loss
-        loss.backward()
-        optimizer.step()
-
-        elapsed_time = time.time() - start_time
-        if elapsed_time > 60:  # Dừng trước 60 giây để lưu checkpoint
-            torch.save(model.state_dict(), "./checkpoint/model.pt")
-            return False, "Checkpoint saved. Training paused."
-
-    return True, "Batch training completed."
-
-
-def train_step(file=None):
-    if file:
-        load_data(file)
-
-    start_idx = 0
-    batch_size = 8
-    total_samples = len(global_data)
-
-    while start_idx < total_samples:
-        end_idx = min(start_idx + (batch_size * 10), total_samples)  # Chia nhỏ dữ liệu để xử lý nhanh
-        dataloader = get_dataloader(start_idx, end_idx, batch_size)
-
-        start_time = time.time()
-        success, message = train_batch(dataloader)
-        elapsed_time = time.time() - start_time
-
-        if elapsed_time >= 100:  # Kết thúc trước khi hết 60 giây để lưu checkpoint
-            torch.save(model.state_dict(), "./checkpoint/model.pt")
-            return f"{message}. Training paused after {elapsed_time:.2f}s."
-
-        start_idx = end_idx
-
-    torch.save(model.state_dict(), "./checkpoint/model.pt")
-    return "Training completed and model saved."
-
-
-if __name__ == "__main__":
-    iface = gr.Interface(
-        fn=train_step,
-        inputs=gr.File(label="Upload CSV"),
-        outputs="text"
-    )
-    iface.launch()
+import time
+import torch
+from transformers import BertForSequenceClassification, AdamW
+from torch.utils.data import DataLoader, TensorDataset
+from transformers import BertTokenizer
+import gradio as gr
+import pandas as pd
+import os
+import spaces
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(device)
+
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+model.to(device)
+
+optimizer = AdamW(model.parameters(), lr=1e-5)
+
+global_data = None
+
+def load_data(file):
+    global global_data
+    df = pd.read_csv(file)
+    inputs = tokenizer(df['text'].tolist(), padding=True, truncation=True, return_tensors="pt")  # Mã hóa văn bản
+    labels = torch.tensor(df['lable'].tolist()).long()  #
+    global_data = TensorDataset(inputs['input_ids'], inputs['attention_mask'], labels)
+
+    print(global_data)
+
+def get_dataloader(start, end, batch_size=8):
+    global global_data
+    subset = torch.utils.data.Subset(global_data, range(start, end))
+    return DataLoader(subset, batch_size=batch_size)
+
+@spaces.GPU(duration=120)
+def train_batch(dataloader):
+    model.train()
+    start_time = time.time()
+
+    for step, batch in enumerate(dataloader):
+        input_ids, attention_mask, labels = batch
+        input_ids, attention_mask, labels = input_ids.to(device), attention_mask.to(device), labels.to(device)
+
+        optimizer.zero_grad()
+        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
+        loss = outputs.loss
+        loss.backward()
+        optimizer.step()
+
+        elapsed_time = time.time() - start_time
+        if elapsed_time > 10:  # Dừng trước 60 giây để lưu checkpoint
+            print("save checkpoint")
+            torch.save(model.state_dict(), "./checkpoint/model.pt")
+            return False, "Checkpoint saved. Training paused."
+
+    return True, "Batch training completed."
+
+
+def train_step(file=None):
+    if file:
+        load_data(file)
+
+    start_idx = 0
+    batch_size = 8
+    total_samples = len(global_data)
+
+    while start_idx < total_samples:
+        print(start_idx)
+        end_idx = min(start_idx + (batch_size * 10), total_samples)  # Chia nhỏ dữ liệu để xử lý nhanh
+        dataloader = get_dataloader(start_idx, end_idx, batch_size)
+
+        start_time = time.time()
+        success, message = train_batch(dataloader)
+        elapsed_time = time.time() - start_time
+
+        if elapsed_time >= 10:  # Kết thúc trước khi hết 60 giây để lưu checkpoint
+            torch.save(model.state_dict(), "./checkpoint/model.pt")
+            return f"{message}. Training paused after {elapsed_time:.2f}s."
+
+        start_idx = end_idx
+
+    torch.save(model.state_dict(), "./checkpoint/model.pt")
+    return "Training completed and model saved."
+
+
+if __name__ == "__main__":
+    iface = gr.Interface(
+        fn=train_step,
+        inputs=gr.File(label="Upload CSV"),
+        outputs="text"
+    )
+    iface.launch()