Avoiding re-loading already loaded models. Stated unload functionality as not-implemented.

detect-pretrain-code-contamination/src/run.py

@@ -23,6 +23,8 @@ import sys
 import gc
 import pickle
 
+models = {}
+
 def save_data(filename, data):
     with open(filename, 'wb') as filehandle:
         # store the data as binary data stream
@@ -36,20 +38,18 @@ def load_data(filename):
     return loaded_data
 
 def unload_model(model,tokenizer):
-    model = model.cpu()
-    del model
-    del tokenizer
-    time.sleep(0.5)
-    gc.collect()
-    torch.cuda.empty_cache() 
+    print("[X] Cannot unload model! Functionality not implemented!")
 
 def load_model(name1):
-    model1 = AutoModelForCausalLM.from_pretrained(name1, return_dict=True, device_map='auto')
-    model1.eval()
-    tokenizer1 = AutoTokenizer.from_pretrained(name1)
+    if name1 not in models:
+        model1 = AutoModelForCausalLM.from_pretrained(name1, return_dict=True, device_map='auto')
+        model1.eval()
+        tokenizer1 = AutoTokenizer.from_pretrained(name1)
 
-    tokenizer1.pad_token = tokenizer1.eos_token 
-    return model1, tokenizer1
+        tokenizer1.pad_token = tokenizer1.eos_token
+        models[name1] = model1
+        models[name1 + "_tokenizer"] = tokenizer1
+    return models[name1], models[name1 + "_tokenizer"]
 
 def calculatePerplexity(sentence, model, tokenizer, gpu):
     """
@@ -105,6 +105,7 @@ def get_neighbors(text,ref_loss,model2,tokenizer2,ratio_gen):
     return neighbors_dl
 
 def evaluate_data(test_data, col_name, target_model, ref_model, ratio_gen, data_name):
+    global model1,model2,tokenizer1,tokenizer2
     print(f"all data size: {len(test_data)}")
     random.seed(0)
     random.shuffle(test_data)