Ensire works on CUDA for extra speed

supporting/RAGAS_Service/train_roberta.py

@@ -113,6 +113,11 @@ def compute_metrics(eval_pred):
         "recall": recall_score(labels, preds, average="weighted", zero_division=0),
     }
 
+def main():
+    torch.multiprocessing.set_start_method('fork')
+    print("CUDA available:", torch.cuda.is_available())
+    print("CUDA device count:", torch.cuda.device_count())
+    print("Current device:", torch.cuda.current_device() if torch.cuda.is_available() else "CPU")
     texts, labels = load_dataset_from_csv("../../data/classify.csv")
 
     tokenizer = RobertaTokenizer.from_pretrained(model_name)
@@ -189,7 +194,8 @@ training_args = TrainingArguments(
         save_strategy="epoch",
         metric_for_best_model="f1",
         greater_is_better=True,
-    dataloader_pin_memory=False
+        dataloader_num_workers=4,
+        dataloader_pin_memory=True
     )
 
     train_dataset = TextDataset(train_encodings, train_labels)
@@ -214,3 +220,8 @@ for k, v in metrics.items():
 
     trainer.save_model("./roberta_classifier")
     tokenizer.save_pretrained("./roberta_classifier")
+
+
+
+if __name__ == "__main__":
+    main()