humanist96/FinGPT_Forecaster

1import re2import os3import datasets4from sklearn.metrics import accuracy_score, mean_squared_error5from collections import defaultdict6from rouge_score import rouge_scorer7 8 9lora_module_dict = {10    'chatglm2': ['query_key_value'],11    'llama2': [12        'q_proj', 'k_proj', 'v_proj',13        'o_proj', 'gate_proj', 'up_proj', 'down_proj',14        # 'embed_tokens', 'lm_head',15    ],16}17 18 19def tokenize(args, tokenizer, feature):20    21    prompt_ids = tokenizer.encode(22        feature['prompt'].strip(), padding=False,23        max_length=args.max_length, truncation=True24    )25    26    target_ids = tokenizer.encode(27        feature['answer'].strip(), padding=False,28        max_length=args.max_length, truncation=True, add_special_tokens=False29    )30    31    input_ids = prompt_ids + target_ids32    exceed_max_length = len(input_ids) >= args.max_length33    34     # Add EOS Token35    if input_ids[-1] != tokenizer.eos_token_id and not exceed_max_length:36        input_ids.append(tokenizer.eos_token_id)37    38    label_ids = [tokenizer.pad_token_id] * len(prompt_ids) + input_ids[len(prompt_ids):]39    40    return {41        "input_ids": input_ids,42        "labels": label_ids,43        "exceed_max_length": exceed_max_length44    }45 46 47def parse_model_name(name, from_remote=False):48    49    if name == 'chatglm2':50        return 'THUDM/chatglm2-6b' if from_remote else 'base_models/chatglm2-6b'51    elif name == 'llama2':52        return 'meta-llama/Llama-2-7b-chat-hf' if from_remote else 'base_models/Llama-2-7b-chat-hf'53    else:54        raise ValueError(f"Undefined base model {name}")55        56    57def load_dataset(names, from_remote=False):58    59    dataset_names = [d for d in names.split(',')]60    dataset_list = []61    62    for name in dataset_names:63        rep = 164        if not os.path.exists(name):65            rep = int(name.split('*')[1]) if '*' in name else 166            name = ('FinGPT/fingpt-forecaster-' if from_remote else 'data/fingpt-forecaster-') + name.split('*')[0]67        tmp_dataset = datasets.load_dataset(name) if from_remote else datasets.load_from_disk(name)68    69        if 'test' not in tmp_dataset:70            tmp_dataset = tmp_dataset.train_test_split(0.2, shuffle=True, seed=42)   71        dataset_list.extend([tmp_dataset] * rep)72    73    return dataset_list74 75 76def parse_answer(answer):77    78    match_res = re.match(r"^\s*\[Positive Developments\]:\s*(.*)\s*\[Potential Concerns\]:\s*(.*)\s*\[Prediction & Analysis\]:\s*(.*)\s*$", answer, flags=re.DOTALL)79    if not match_res:80        return None81    82    pros, cons, pna = match_res.group(1), match_res.group(2), match_res.group(3)83        84    match_res = re.match(r'^Prediction:\s*(.*)\s*Analysis:\s*(.*)\s*$', pna, flags=re.DOTALL)85    if not match_res:86        return None87        88    pred, anal = match_res.group(1), match_res.group(2)89        90    if re.search(r'up|increase', pred.lower()):91        pred_bin = 192    elif re.search(r'down|decrease|decline', pred.lower()):93        pred_bin = -194    else:95        pred_bin = 096            97    match_res = re.search(r'(\d)-(\d)%', pred)98    if not match_res:99        match_res = re.search(r'(?:more than )?(\d)+?%', pred)    100        101    pred_margin = pred_bin * (int(match_res.group(1)) + 0.5) if match_res else 0.102        103    return {104        "positive developments": pros,105        "potential concerns": cons,106        "prediction": pred_margin,107        "prediction_binary": pred_bin,108        "analysis": anal109    }110    111 112def calc_rouge_score(references, answers):113    114    scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'], use_stemmer=True)115        116    scores_per_pair = [scorer.score(ref, ans) for ref, ans in zip(references, answers)]117    118    rouge1 = sum(score['rouge1'].fmeasure for score in scores_per_pair) / len(scores_per_pair)119    rouge2 = sum(score['rouge2'].fmeasure for score in scores_per_pair) / len(scores_per_pair)120    rougeL = sum(score['rougeL'].fmeasure for score in scores_per_pair) / len(scores_per_pair)121    122    return {'rouge1': rouge1, 'rouge2': rouge2, 'rougeL': rougeL}123 124    125def calc_metrics(answers, gts):126    127    answers_dict = defaultdict(list)128    gts_dict = defaultdict(list)129    130    for answer, gt in zip(answers, gts):131        answer_dict = parse_answer(answer)132        gt_dict = parse_answer(gt)133        134        if answer_dict and gt_dict:135            for k in answer_dict.keys():136                answers_dict[k].append(answer_dict[k])137                gts_dict[k].append(gt_dict[k])138    139    if not answers_dict['prediction']:140        return {}141    142    bin_acc = accuracy_score(gts_dict['prediction_binary'], answers_dict['prediction_binary'])143    mse = mean_squared_error(gts_dict['prediction'], answers_dict['prediction'])144    145    pros_rouge_scores = calc_rouge_score(gts_dict['positive developments'], answers_dict['positive developments'])146    cons_rouge_scores = calc_rouge_score(gts_dict['potential concerns'], answers_dict['potential concerns'])147    anal_rouge_scores = calc_rouge_score(gts_dict['analysis'], answers_dict['analysis'])148                              149    print(f"\nBinary Accuracy: {bin_acc:.2f}  |  Mean Square Error: {mse:.2f}")150    print(f"\nRouge Score of Positive Developments: {pros_rouge_scores}")151    print(f"\nRouge Score of Potential Concerns: {cons_rouge_scores}")152    print(f"\nRouge Score of Summary Analysis: {anal_rouge_scores}")153                              154    return {155        "valid_count": len(answers_dict['prediction']),156        "bin_acc": bin_acc,157        "mse": mse,158        "pros_rouge_scores": pros_rouge_scores,159        "cons_rouge_scores": cons_rouge_scores,160        "anal_rouge_scores": anal_rouge_scores161    }162