Add threads

app.py

@@ -49,8 +49,8 @@ st.write('Loading the pretrained model ...')
 model_name = "CarolXia/pii-kd-deberta-v2"
 # config = PeftConfig.from_pretrained(model_name)
 model = DebertaV2ForTokenClassification.from_pretrained(model_name, token=st.secrets["HUGGINGFACE_TOKEN"])
-# if torch.cuda.is_available():
-#     model = model.to("cuda")
+if torch.cuda.is_available():
+    model = model.to("cuda")
 # Try quantization instead
 # model = AutoModelForTokenClassification.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
 tokenizer = DebertaV2Tokenizer.from_pretrained("microsoft/mdeberta-v3-base", token=st.secrets["HUGGINGFACE_TOKEN"])
@@ -141,18 +141,27 @@ entity_set=dict()
 dataset = load_dataset("Isotonic/pii-masking-200k", split="train")
 unmasked_text = dataset['unmasked_text'] # This will load the entire column inmemory. Must do this to avoid I/O delay later
 
-st.write('Size of the dataset ', dataset.num_rows)
-sizes = [0] * 2
+st.write('Number of rows in the dataset ', dataset.num_rows)
+sizes = [0] * 5
 start = time.time()
-t1 = threading.Thread(target=process_datasets, args=(0, 25, unmasked_text, sizes, 0, entity_set, []))
-t2 = threading.Thread(target=process_datasets, args=(25, 50, unmasked_text, sizes, 1, entity_set, []))
-with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], profile_memory=True, record_shapes=True) as prof:
-    process_datasets(0, 50, unmasked_text, sizes, 0, entity_set, [])
-    # t1.start()
-    # t2.start()
-
-    # t1.join()
-    # t2.join()
+# t0 = threading.Thread(target=process_datasets, args=(0, 50, unmasked_text, sizes, 0, entity_set, []))
+# t1 = threading.Thread(target=process_datasets, args=(25, 50, unmasked_text, sizes, 1, entity_set, []))
+# t2 = threading.Thread(target=process_datasets, args=(20, 30, unmasked_text, sizes, 2, entity_set, []))
+# t3 = threading.Thread(target=process_datasets, args=(30, 40, unmasked_text, sizes, 3, entity_set, []))
+# t4 = threading.Thread(target=process_datasets, args=(40, 50, unmasked_text, sizes, 4, entity_set, []))
+# with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], profile_memory=True, record_shapes=True) as prof:
+process_datasets(0, 50, unmasked_text, sizes, 0, entity_set, [])
+# t0.start()
+# t1.start()
+# t2.start()
+# t3.start()
+# t4.start()
+
+# t0.join()
+# t1.join()
+# t2.join()
+# t3.join()
+# t4.join()
 
 end = time.time()
 length = end - start