Upload run.py

run.py

@@ -17,8 +17,6 @@ def tp_tf_test(test_dataset, metric_selector, model_selector, queries_selector,
 	queries_dataset_path = queries_selector.replace('-'+queries_data_file, '')
 	queries_dataset_split = {'queries': queries_data_file}
 	queries_dataset = load_dataset(queries_dataset_path,  data_files = queries_dataset_split)['queries']
-	#queries_data_files = {'queries': queries_selector}
-	#queries_dataset = load_dataset('gorkaartola/SDG_queries', data_files = queries_data_files)['queries']
 	
 	#Load prompt_________________________________
 	prompt = prompt_selector
@@ -89,10 +87,11 @@ def tp_tf_test(test_dataset, metric_selector, model_selector, queries_selector,
 	results_test = pd.read_csv('Reports/ZS inference tables/ZS-inference-table_Model-' + op.models[model_selector] + '_Queries-' + op.queries[queries_dataset_path][queries_data_file]  + '_Prompt-' + op.prompts[prompt_selector] + '.csv')
 	'''
 	#Analize predictions_______________________________
-	def logits_labels(raw): 
+	def logits_labels(raw):
+		classes = raw["dataset_labels"].unique()
 		raw_logits = raw.iloc[:,2:]
-		logits = np.zeros(shape=(len(raw_logits.index),17))
-		for i in range(17):
+		logits = np.zeros(shape=(len(raw_logits.index),len(classes)))
+		for i in range(len(classes)):
 			queries = queries_dataset.filter(lambda x: x['label_ids'] == i)['SDGquery']
 			logits[:,i]=raw_logits[queries].max(axis=1)
 		labels = raw[["dataset_labels","nli_labels"]]