Spaces:
Running
Running
| from sklearn.metrics import accuracy_score, f1_score, classification_report | |
| from datasets import load_dataset, load_from_disk | |
| from tqdm import tqdm | |
| import datasets | |
| import torch | |
| from torch.utils.data import DataLoader | |
| from functools import partial | |
| from pathlib import Path | |
| from fingpt.FinGPT_Benchmark.utils import * | |
| import sys | |
| sys.path.append('../') | |
| def binary2multi(dataset): | |
| pred, label = [], [] | |
| tmp_pred, tmp_label = [], [] | |
| for i, row in dataset.iterrows(): | |
| tmp_pred.append(row['pred']) | |
| tmp_label.append(row['label']) | |
| if (i + 1) % 9 == 0: | |
| pred.append(tmp_pred) | |
| label.append(tmp_label) | |
| tmp_pred, tmp_label = [], [] | |
| return pred, label | |
| def map_output(feature): | |
| pred = 1 if 'yes' in feature['out_text'].lower() else 0 | |
| label = 1 if 'yes' in feature['output'].lower() else 0 | |
| return {'label': label, 'pred': pred} | |
| def test_headline(args, model, tokenizer): | |
| # dataset = load_from_disk('../data/fingpt-headline')['test'] | |
| dataset = load_from_disk(Path(__file__).parent.parent / 'data/fingpt-headline-instruct')['test'] | |
| dataset = dataset.map(partial(test_mapping, args), load_from_cache_file=False) | |
| def collate_fn(batch): | |
| inputs = tokenizer( | |
| [f["prompt"] for f in batch], return_tensors='pt', | |
| padding=True, max_length=args.max_length, | |
| return_token_type_ids=False | |
| ) | |
| return inputs | |
| dataloader = DataLoader(dataset, batch_size=args.batch_size, collate_fn=collate_fn, shuffle=False) | |
| out_text_list = [] | |
| log_interval = len(dataloader) // 5 | |
| for idx, inputs in enumerate(tqdm(dataloader)): | |
| inputs = {key: value.to(model.device) for key, value in inputs.items()} | |
| res = model.generate(**inputs, max_length=args.max_length, eos_token_id=tokenizer.eos_token_id) | |
| res_sentences = [tokenizer.decode(i, skip_special_tokens=True) for i in res] | |
| tqdm.write(f'{idx}: {res_sentences[0]}') | |
| if (idx + 1) % log_interval == 0: | |
| tqdm.write(f'{idx}: {res_sentences[0]}') | |
| out_text = [o.split("Answer: ")[1] for o in res_sentences] | |
| out_text_list += out_text | |
| torch.cuda.empty_cache() | |
| dataset = dataset.add_column("out_text", out_text_list) | |
| dataset = dataset.map(map_output, load_from_cache_file=False) | |
| dataset = dataset.to_pandas() | |
| print(dataset) | |
| dataset.to_csv('tmp.csv') | |
| # binary | |
| acc = accuracy_score(dataset["label"], dataset["pred"]) | |
| f1 = f1_score(dataset["label"], dataset["pred"], average="binary") | |
| # multi-class | |
| pred, label = binary2multi(dataset) | |
| print(f"\n|| Acc: {acc} || F1 binary: {f1} ||\n") | |
| print(classification_report(label, pred, digits=4, target_names=['price or not', 'price up', 'price stable', | |
| 'price down', 'price past', 'price future', | |
| 'event past', 'event future', 'asset comp'])) | |
| return dataset |