Spaces:
Runtime error
Runtime error
| from transformers import AutoTokenizer, AutoModelForSequenceClassification, DataCollatorWithPadding | |
| from datasets import load_dataset, load_metric | |
| import evaluate | |
| from torch.utils.data import DataLoader | |
| import torch | |
| import numpy as np | |
| import pandas as pd | |
| import options as op | |
| def tp_tf_test(model_selector, test_dataset, queries_selector, prompt_selector, metric_selector, prediction_strategy_selector): | |
| #Load test dataset___________________________ | |
| test_dataset = load_dataset(test_dataset)['test'] | |
| #Load queries________________________________ | |
| queries_data_file = queries_selector.split('-')[-1] | |
| 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'] | |
| #Load prompt_________________________________ | |
| prompt = prompt_selector | |
| #Load prediction strategias__________________ | |
| prediction_strategies = prediction_strategy_selector | |
| #Load model, tokenizer and collator__________ | |
| model = AutoModelForSequenceClassification.from_pretrained(model_selector) | |
| if torch.cuda.is_available(): | |
| device = torch.device("cuda") | |
| model.to(device) | |
| tokenizer = AutoTokenizer.from_pretrained(model_selector) | |
| data_collator = DataCollatorWithPadding(tokenizer=tokenizer) | |
| #Calculate and save predictions______________ | |
| #''' | |
| def tokenize_function(example, prompt = '', query = ''): | |
| queries = [] | |
| for i in range(len(example['title'])): | |
| queries.append(prompt + query) | |
| tokenize = tokenizer(example['title'], queries, truncation='only_first') | |
| #tokenize['query'] = queries | |
| return tokenize | |
| results_test = pd.DataFrame() | |
| for query_data in queries_dataset: | |
| query = query_data['query'] | |
| tokenized_test_dataset = test_dataset.map(tokenize_function, batched = True, fn_kwargs = {'prompt' : prompt, 'query' : query}) | |
| columns_to_remove = test_dataset.column_names | |
| for column_name in ['label_ids', 'nli_label']: | |
| columns_to_remove.remove(column_name) | |
| tokenized_test_dataset_for_inference = tokenized_test_dataset.remove_columns(columns_to_remove) | |
| tokenized_test_dataset_for_inference.set_format('torch') | |
| dataloader = DataLoader( | |
| tokenized_test_dataset_for_inference, | |
| batch_size=8, | |
| collate_fn = data_collator, | |
| ) | |
| values = [] | |
| labels = [] | |
| nli_labels =[] | |
| for batch in dataloader: | |
| if torch.cuda.is_available(): | |
| data = {k: v.to(device) for k, v in batch.items() if k not in ['labels', 'nli_label']} | |
| else: | |
| data = {k: v for k, v in batch.items() if k not in ['labels', 'nli_label']} | |
| with torch.no_grad(): | |
| outputs = model(**data) | |
| logits = outputs.logits | |
| entail_contradiction_logits = logits[:,[0,2]] | |
| probs = entail_contradiction_logits.softmax(dim=1) | |
| predictions = probs[:,1].tolist() | |
| label_ids = batch['labels'].tolist() | |
| nli_label_ids = batch['nli_label'].tolist() | |
| for prediction, label, nli_label in zip(predictions, label_ids, nli_label_ids): | |
| values.append(prediction) | |
| labels.append(label) | |
| nli_labels.append(nli_label) | |
| results_test['dataset_labels'] = labels | |
| results_test['nli_labels'] = nli_labels | |
| results_test[query] = values | |
| results_test.to_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', index = False) | |
| #''' | |
| #Load saved predictions____________________________ | |
| ''' | |
| 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): | |
| classes = raw["dataset_labels"].unique() | |
| raw_logits = raw.iloc[:,2:] | |
| 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)['query'] | |
| logits[:,i]=raw_logits[queries].max(axis=1) | |
| labels = raw[["dataset_labels","nli_labels"]] | |
| labels = np.array(labels).astype(int) | |
| return logits, labels | |
| predictions, references = logits_labels(results_test) | |
| prediction_strategies = [op.prediction_strategy_options[x] for x in prediction_strategy_selector] | |
| metric = evaluate.load(metric_selector) | |
| metric.add_batch(predictions = predictions, references = references) | |
| results = metric.compute(prediction_strategies = prediction_strategies) | |
| prediction_strategies_names = '-'.join(prediction_strategy_selector).replace(" ", "") | |
| output_filename = 'Reports/report-Model-' + op.models[model_selector] + '_Queries-' + op.queries[queries_dataset_path][queries_data_file] + '_Prompt-' + op.prompts[prompt_selector] + '_Strategies-'+ prediction_strategies_names +'.csv' | |
| with open(output_filename, 'a') as results_file: | |
| for result in results: | |
| results[result].to_csv(results_file, mode='a', index_label = result) | |
| print(results[result], '\n') | |
| return output_filename | |