This repository introduces MentaLLaMA, the first open-source instruction following large language model for interpretable mental health analysis.
import os
import argparse
import pandas as pd

import torch
from transformers import AutoModel, AutoTokenizer, LlamaForCausalLM, LlamaTokenizer

from sklearn.metrics import f1_score, confusion_matrix, accuracy_score, classification_report, \
    precision_recall_fscore_support, precision_score, recall_score

def load_instruction_test_data():
    test_data = {}
    for root, ds, fs in os.walk("../test_data/test_instruction"):
        for fn in fs:
            data = pd.read_csv(os.path.join(root, fn))
            texts = data['query'].to_list()
            labels = data['gpt-3.5-turbo'].to_list()
            test_data[fn.split('.')[0]] = [texts, labels]
    return test_data

def load_complete_test_data():
    test_data = {}
    for root, ds, fs in os.walk("../test_data/test_complete"):
        for fn in fs:
            data = pd.read_csv(os.path.join(root, fn))
            texts = data['query'].to_list()
            labels = data['gpt-3.5-turbo'].to_list()
            test_data[fn.split('.')[0]] = [texts, labels]
    return test_data

def load_expert_data():
    test_data = {}
    for root, ds, fs in os.walk("../human_evaluation/test_instruction_expert"):
        for fn in fs:
            data = pd.read_csv(os.path.join(root, fn))
            texts = data['query'].to_list()
            labels = data['gpt-3.5-turbo'].to_list()
            test_data[fn.split('.')[0]] = [texts, labels]
    return test_data

def generate_response(model, tokenizer, test_data, device, batch_size):
    generated_text = {}
    goldens = {}

    model.to(device)

    for dataset_name in test_data.keys():
        #if dataset_name not in ['DR', 'dreaddit']:
        #    continue
        print('Generating for dataset: {}'.format(dataset_name))
        queries, golden = test_data[dataset_name]
        goldens[dataset_name]  = golden
        responses = []

        for i in range(0, len(queries), batch_size):
            batch_data = queries[i: min(i+batch_size, len(queries))]
            #print(batch_data[:2])
            inputs = tokenizer(batch_data, return_tensors="pt", padding=True)
            #print(inputs)
            #final_input = inputs.input_ids
            input_ids = inputs.input_ids.to(device)
            attention_mask = inputs.attention_mask.to(device)
            #print(final_input)
            generate_ids = model.generate(input_ids, attention_mask=attention_mask, max_length=2048)
            for j in range(generate_ids.shape[0]):
                truc_ids = generate_ids[j][len(input_ids[j]) :]
                response = tokenizer.decode(truc_ids, skip_special_tokens=True, spaces_between_special_tokens=False)
                responses.append(response)
            print(i)
        generated_text[dataset_name] = responses

    return generated_text, goldens

def save_output(generated_text, goldens, output_path):
    if not os.path.exists("../model_output/"):
        os.mkdir("../model_output/")
    if not os.path.exists("../model_output/"+output_path):
        os.mkdir("../model_output/"+output_path)
    for dataset_name in generated_text.keys():
        output = {'goldens': goldens[dataset_name], 'generated_text': generated_text[dataset_name]}
        output = pd.DataFrame(output, index=None)
        output.to_csv("{}/{}/{}.csv".format('../model_output',
                                         output_path, dataset_name), index=False, escapechar='\\')

def calculate_f1(generated, goldens, output_path):
    for dataset_name in generated.keys():
        golden = goldens[dataset_name]
        outputs = generated[dataset_name]

        output_label = []
        golden_label = []
        count = 0

        for ref, output in zip(golden, outputs):
            ref_an = ref.split("Reasoning:")[0]
            output_an = output.split("Reasoning:")[0]
            #print(output)
            #output_an = str(output)[:70]

            if dataset_name == 'swmh':
                if 'no mental' in output_an.lower():
                    output_label.append(0)
                elif 'suicide' in output_an.lower():
                    output_label.append(1)
                elif 'depression' in output_an.lower():
                    output_label.append(2)
                elif 'anxiety' in output_an.lower():
                    output_label.append(3)
                elif 'bipolar' in output_an.lower():
                    output_label.append(4)
                else:
                    count += 1
                    output_label.append(0)
                    #print(output)

                if 'no mental' in ref_an.lower():
                    golden_label.append(0)
                elif 'suicide' in ref_an.lower():
                    golden_label.append(1)
                elif 'depression' in ref_an.lower():
                    golden_label.append(2)
                elif 'anxiety' in ref_an.lower():
                    golden_label.append(3)
                elif 'bipolar' in ref_an.lower():
                    golden_label.append(4)
                else:
                    output_label = output_label[:-1]

            elif dataset_name == 't-sid':
                if 'depression' in output_an.lower():
                    output_label.append(2)
                elif 'suicide' in output_an.lower():
                    output_label.append(1)
                elif 'ptsd' in output_an.lower():
                    output_label.append(3)
                elif 'no mental' in output_an.lower():
                    output_label.append(0)
                else:
                    count += 1
                    #print(output)
                    output_label.append(0)
                    #print(output)

                if 'depression' in ref_an.lower():
                    golden_label.append(2)
                elif 'suicide or self-harm' in ref_an.lower():
                    golden_label.append(1)
                elif 'ptsd' in ref_an.lower():
                    golden_label.append(3)
                elif 'no mental disorders' in ref_an.lower():
                    golden_label.append(0)

            elif dataset_name in ['CLP', 'DR', 'dreaddit', 'loneliness', 'Irf', 'MultiWD']:
                if 'yes' in output_an.lower():
                    output_label.append(1)
                elif 'no' in output_an.lower():
                    output_label.append(0)
                else:
                    count += 1
                    output_label.append(0)
                    #print(output)

                if 'yes' in ref_an.lower():
                    golden_label.append(1)
                elif 'no' in ref_an.lower():
                    golden_label.append(0)

            elif dataset_name == 'SAD':
                if 'school' in output_an.lower():
                    output_label.append(0)
                elif 'financial' in output_an.lower():
                    output_label.append(1)
                elif 'family' in output_an.lower():
                    output_label.append(2)
                elif 'social' in output_an.lower():
                    output_label.append(3)
                elif 'work' in output_an.lower():
                    output_label.append(4)
                elif 'health' in output_an.lower():
                    output_label.append(5)
                elif 'emotional' in output_an.lower():
                    output_label.append(6)
                elif 'decision' in output_an.lower():
                    output_label.append(7)
                elif 'other' in output_an.lower():
                    output_label.append(8)
                else:
                    count += 1
                    output_label.append(8)
                    #print(output_an)

                if 'school' in ref_an.lower():
                    golden_label.append(0)
                elif 'financial problem' in ref_an.lower():
                    golden_label.append(1)
                elif 'family issues' in ref_an.lower():
                    golden_label.append(2)
                elif 'social relationships' in ref_an.lower():
                    golden_label.append(3)
                elif 'work' in ref_an.lower():
                    golden_label.append(4)
                elif 'health issues' in ref_an.lower():
                    golden_label.append(5)
                elif 'emotional turmoil' in ref_an.lower():
                    golden_label.append(6)
                elif 'everyday decision making' in ref_an.lower():
                    golden_label.append(7)
                elif 'other causes' in ref_an.lower():
                    golden_label.append(8)

            elif dataset_name == 'CAMS':
                if 'none' in output_an.lower():
                    output_label.append(0)
                elif 'bias' in output_an.lower():
                    output_label.append(1)
                elif 'jobs' in output_an.lower():
                    output_label.append(2)
                elif 'medication' in output_an.lower():
                    output_label.append(3)
                elif 'relationship' in output_an.lower():
                    output_label.append(4)
                elif 'alienation' in output_an.lower():
                    output_label.append(5)
                else:
                    count += 1
                    output_label.append(0)
                    #print(output_an)

                if 'none' in ref_an.lower():
                    golden_label.append(0)
                elif 'bias or abuse' in ref_an.lower():
                    golden_label.append(1)
                elif 'jobs and career' in ref_an.lower():
                    golden_label.append(2)
                elif 'medication' in ref_an.lower():
                    golden_label.append(3)
                elif 'relationship' in ref_an.lower():
                    golden_label.append(4)
                elif 'alienation' in ref_an.lower():
                    golden_label.append(5)

        avg_accuracy = round(accuracy_score(golden_label, output_label) * 100, 2)
        weighted_f1 = round(f1_score(golden_label, output_label, average='weighted') * 100, 2)
        micro_f1 = round(f1_score(golden_label, output_label, average='micro') * 100, 2)
        macro_f1 = round(f1_score(golden_label, output_label, average='macro') * 100, 2)
        # recall = round(recall_score(f_labels, outputs, average='weighted')*100, 2)
        result = "Dataset: {}, average acc:{}, weighted F1 {}, micro F1 {}, macro F1 {}, OOD count: {}\n".format(dataset_name,
                                                                                             avg_accuracy, weighted_f1,
                                                                                             micro_f1, macro_f1, count)
        print(result)
        with open("{}/{}/results.txt".format('../model_output', output_path), 'a+') as f:
            f.write(result)
        print(count)

def main(model_path: str, model_output_path: str, batch_size: int, test_dataset: str, rule_calculate: bool, llama: bool, device: str, cuda: bool):
    if llama:
        model = LlamaForCausalLM.from_pretrained(model_path)
        tokenizer = LlamaTokenizer.from_pretrained(model_path, padding_side='left')
    else:
        model = AutoModel.from_pretrained(model_path)
        tokenizer = AutoTokenizer.from_pretrained(model_path)

    tokenizer.pad_token = tokenizer.unk_token

    if test_dataset == 'IMHI':
        test_data = load_instruction_test_data()
    elif test_dataset == 'IMHI-completion':
        test_data = load_complete_test_data()
    elif test_dataset == 'expert':
        test_data = load_expert_data()
    generated_text, goldens = generate_response(model, tokenizer, test_data, device, batch_size)
    save_output(generated_text, goldens, model_output_path)
    if rule_calculate:
        calculate_f1(generated_text, goldens, model_output_path)



if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='The IMHI benchmark.')
    parser.add_argument('--model_path', type=str, default='')
    parser.add_argument('--model_output_path', type=str)
    parser.add_argument('--cuda', action='store_true')
    parser.add_argument('--batch_size', type=int, default=24)
    parser.add_argument('--test_dataset', type=str, choices=['IMHI', 'IMHI-completion', 'expert'])
    parser.add_argument('--rule_calculate', action='store_true')
    parser.add_argument('--llama', action='store_true')

    args = parser.parse_args()
    args = vars(args)
    device = torch.device("cuda:0" if torch.cuda.is_available() and args['cuda'] is True else "cpu")
    args['device'] = device

    main(**args)
src/IMHI.py
SteveKGYang/MentalLLaMA
from transformers import BertTokenizer, BertForSequenceClassification
import numpy as np
import argparse
import pandas as pd
import torch
import os

from sklearn.metrics import f1_score, confusion_matrix, accuracy_score, classification_report, \
    precision_recall_fscore_support, precision_score, recall_score

def get_dict(dataset_name):
    if dataset_name == 'CAMS':
        labels = {0: 'none', 1: 'bias', 2: 'job', 3: 'medication', 4: 'relation', 5: 'alienation'}
        num_labels = 6
    elif dataset_name == 'CLP':
        labels = {0: 'no', 1: 'yes'}
        num_labels = 2
    elif dataset_name == 'DR':
        labels = {0: 'no', 1: 'yes'}
        num_labels = 2
    elif dataset_name == 'dreaddit':
        labels = {0: 'no', 1: 'yes'}
        num_labels = 2
    elif dataset_name == 'Irf':
        labels = {0: 'no', 1: 'yes'}
        num_labels = 2
    elif dataset_name == 'loneliness':
        labels = {0: 'no', 1: 'yes'}
        num_labels = 2
    elif dataset_name == 'MultiWD':
        labels = {0: 'no', 1: 'yes'}
        num_labels = 2
    elif dataset_name == 'SAD':
        labels = {0: 'school', 1: 'financial',2: 'family', 3: 'social',4: 'work', 5: 'health',6: 'emotional', 7: 'everyday',8: 'other'}
        num_labels = 9
    elif dataset_name == 'swmh':
        labels = {0: 'depression', 1: 'suicide', 2: 'anxiety', 3: 'bipolar', 4: 'no mental'}
        num_labels = 5
    elif dataset_name == 't-sid':
        labels = {0: 'depression', 1: 'suicide', 2: 'ptsd', 3: 'control'}
        num_labels = 4
    else:
        raise Exception("ERROR! please choose the correct dataset!")

    return labels, num_labels

def calculate_f1(goldens, final_labels, dataset_name):
    golden_label = []
    output_label = []
    for golden, label in zip(goldens, final_labels):
        ref_an = golden.split("Reasoning:")[0]
        output_an = label.strip()

        if 'swmh' in dataset_name:
            if 'no mental' in output_an.lower():
                output_label.append(0)
            elif 'suicide' in output_an.lower():
                output_label.append(1)
            elif 'depression' in output_an.lower():
                output_label.append(2)
            elif 'anxiety' in output_an.lower():
                output_label.append(3)
            elif 'bipolar' in output_an.lower():
                output_label.append(4)
            else:
                raise Exception('Wrong label in predictions for {}'.format(dataset_name))

            if 'no mental' in ref_an.lower():
                golden_label.append(0)
            elif 'suicide' in ref_an.lower():
                golden_label.append(1)
            elif 'depression' in ref_an.lower():
                golden_label.append(2)
            elif 'anxiety' in ref_an.lower():
                golden_label.append(3)
            elif 'bipolar' in ref_an.lower():
                golden_label.append(4)
            else:
                output_label = output_label[:-1]

        elif dataset_name == 't-sid':
            if 'depression' in output_an.lower():
                output_label.append(2)
            elif 'suicide' in output_an.lower():
                output_label.append(1)
            elif 'ptsd' in output_an.lower():
                output_label.append(3)
            elif 'control' in output_an.lower():
                output_label.append(0)
            else:
                raise Exception('Wrong label in predictions for {}'.format(dataset_name))

            if 'depression' in ref_an.lower():
                golden_label.append(2)
            elif 'suicide or self-harm' in ref_an.lower():
                golden_label.append(1)
            elif 'ptsd' in ref_an.lower():
                golden_label.append(3)
            elif 'no mental' in ref_an.lower():
                golden_label.append(0)

        elif dataset_name in ['CLP', 'DR', 'dreaddit', 'loneliness', 'Irf', 'MultiWD']:
            if 'yes' in output_an.lower():
                output_label.append(1)
            elif 'no' in output_an.lower():
                output_label.append(0)
            else:
                raise Exception('Wrong label in predictions for {}'.format(dataset_name))

            if 'yes' in ref_an.lower():
                golden_label.append(1)
            elif 'no' in ref_an.lower():
                golden_label.append(0)

        elif dataset_name == 'SAD':
            if 'school' in output_an.lower():
                output_label.append(0)
            elif 'financial' in output_an.lower():
                output_label.append(1)
            elif 'family' in output_an.lower():
                output_label.append(2)
            elif 'social' in output_an.lower():
                output_label.append(3)
            elif 'work' in output_an.lower():
                output_label.append(4)
            elif 'health' in output_an.lower():
                output_label.append(5)
            elif 'emotional' in output_an.lower():
                output_label.append(6)
            elif 'everyday' in output_an.lower():
                output_label.append(7)
            elif 'other' in output_an.lower():
                output_label.append(8)
            else:
                raise Exception('Wrong label in predictions for {}'.format(dataset_name))

            if 'school' in ref_an.lower():
                golden_label.append(0)
            elif 'financial problem' in ref_an.lower():
                golden_label.append(1)
            elif 'family issues' in ref_an.lower():
                golden_label.append(2)
            elif 'social relationships' in ref_an.lower():
                golden_label.append(3)
            elif 'work' in ref_an.lower():
                golden_label.append(4)
            elif 'health issues' in ref_an.lower():
                golden_label.append(5)
            elif 'emotion turmoil' in ref_an.lower():
                golden_label.append(6)
            elif 'everyday decision making' in ref_an.lower():
                golden_label.append(7)
            elif 'other' in ref_an.lower():
                golden_label.append(8)

        elif dataset_name == 'CAMS':
            if 'none' in output_an.lower():
                output_label.append(0)
            elif 'bias' in output_an.lower():
                output_label.append(1)
            elif 'job' in output_an.lower():
                output_label.append(2)
            elif 'medication' in output_an.lower():
                output_label.append(3)
            elif 'relation' in output_an.lower():
                output_label.append(4)
            elif 'alienation' in output_an.lower():
                output_label.append(5)
            else:
                raise Exception('Wrong label in predictions for {}'.format(dataset_name))

            if 'no causes' in ref_an.lower():
                golden_label.append(0)
            elif 'bias or abuse' in ref_an.lower():
                golden_label.append(1)
            elif 'jobs and career' in ref_an.lower():
                golden_label.append(2)
            elif 'medication' in ref_an.lower():
                golden_label.append(3)
            elif 'relationship' in ref_an.lower():
                golden_label.append(4)
            elif 'alienation' in ref_an.lower():
                golden_label.append(5)
    avg_accuracy = round(accuracy_score(golden_label, output_label) * 100, 2)
    weighted_f1 = round(f1_score(golden_label, output_label, average='weighted') * 100, 2)
    micro_f1 = round(f1_score(golden_label, output_label, average='micro') * 100, 2)
    macro_f1 = round(f1_score(golden_label, output_label, average='macro') * 100, 2)
    print("Dataset: {}, average acc:{}, weighted F1 {}, micro F1 {}, macro F1 {}".format(dataset_name,
                                                                                         avg_accuracy, weighted_f1,
                                                                                         micro_f1, macro_f1))
def main(model_path: str, data_path: str, data_output_path: str, batch_size: int, device: str, cuda: bool, calculate: bool):
    cal_data = {}
    if not os.path.exists(data_output_path):
        os.mkdir(data_output_path)
    for root, ds, fs in os.walk(data_path):
        for fn in fs:
            data = pd.read_csv(os.path.join(root, fn))
            goldens = data['goldens'].to_list()
            generated_text = data['generated_text'].to_list()

            dataset_name = fn.split('.')[0]
            labels, num_labels = get_dict(dataset_name)
            print('Generating for {} dataset.'.format(dataset_name))

            mentalbert = BertForSequenceClassification.from_pretrained(os.path.join(model_path, dataset_name), num_labels=num_labels).to(device)
            tokenizer = BertTokenizer.from_pretrained(os.path.join(model_path, dataset_name))

            all_labels = []
            for i in range(0, len(generated_text), batch_size):
                batch_data = generated_text[i: min(i + batch_size, len(generated_text))]
                inputs = tokenizer(batch_data, return_tensors="pt", padding=True, truncation=True, max_length=512).to(
                    device)
                outputs = mentalbert(**inputs)[0]
                outputs = outputs.cpu().detach().numpy()
                for j in range(0, len(batch_data)):
                    all_labels.append(np.argmax(outputs[j]))
            final_labels = []
            for num in all_labels:
                final_labels.append(labels[num])

            output = {'goldens': goldens, 'generated_text': generated_text, 'labels': final_labels}
            output_file = pd.DataFrame(output, index=None)
            output_file.to_csv(os.path.join(data_output_path, '{}.csv'.format(dataset_name)), index=False, escapechar='\\')

            if calculate:
                cal_data[dataset_name] = [goldens, final_labels]
    if calculate:
        for dataset_name in cal_data.keys():
            goldens, final_labels = cal_data[dataset_name]
            calculate_f1(goldens, final_labels, dataset_name)



if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='The IMHI benchmark.')
    parser.add_argument('--model_path', type=str)
    parser.add_argument('--data_path', type=str, default='')
    parser.add_argument('--data_output_path', type=str)
    parser.add_argument('--cuda', action='store_true')
    parser.add_argument('--calculate', action='store_true')
    parser.add_argument('--batch_size', type=int, default=8)
    args = parser.parse_args()
    args = vars(args)
    device = torch.device("cuda:0" if torch.cuda.is_available() and args['cuda'] is True else "cpu")
    args['device'] = device

    main(**args)
src/label_inference.py
SteveKGYang/MentalLLaMA
import os
import pandas as pd
import evaluate
import pickle
import argparse
import torch

from BARTScore.bart_score import BARTScorer
from GPTScore.gpt3_score import gpt3score

def rouge(gen_dir_name):
    rouge = evaluate.load('rouge')
    score_results = {}

    for root, ds, fs in os.walk("../model_output/{}".format(gen_dir_name)):
        for fn in fs:
            data = pd.read_csv(os.path.join(root, fn))
            dname = fn.split('.')[0]
            predictions = data['generated_text'].to_list()
            references = data['goldens'].to_list()

            result = rouge.compute(predictions=predictions, references=references)
            score_results[dname] = [result['rouge1'], result['rouge2'], result['rougeL']]
            print('Results for {} dataset: {}'.format(dname, score_results[dname]))
    pickle.dump(score_results,
                open('../quality_evaluation_results/rouge_score_{}.pkl'.format(gen_dir_name), 'wb+'))
    return score_results

def bleu(gen_dir_name):
    score_results = {}
    rouge = evaluate.load('bleu')

    for root, ds, fs in os.walk("../model_output/{}".format(gen_dir_name)):
        for fn in fs:
            data = pd.read_csv(os.path.join(root, fn))
            dname = fn.split('.')[0]
            predictions = data['generated_text'].to_list()
            references = data['goldens'].to_list()

            result = rouge.compute(predictions=predictions, references=references)
            score_results[dname] = result['bleu']
            print('Results for {} dataset: {}'.format(dname, score_results[dname]))
    pickle.dump(score_results,
                open('../quality_evaluation_results/bleu_score_{}.pkl'.format(gen_dir_name), 'wb+'))
    return score_results

def GPTScore(gen_dir_name):
    GPT_model = input('Which GPT-based model will you use?')
    api_key = input('Your api key: ')
    score_results = {}
    results = []
    for root, ds, fs in os.walk("../model_output/{}".format(gen_dir_name)):
        for fn in fs:
            data = pd.read_csv(os.path.join(root, fn))
            dname = fn.split('.')[0]
            predictions = data['generated_text'].to_list()
            references = data['goldens'].to_list()
            for prediction, reference in zip(predictions, references):
                score = gpt3score(reference, prediction, gpt3model=GPT_model, api_key=api_key)
                results.append(score)
            score_results[dname] = sum(results) / len(results)
            print('Results for {} dataset: {}'.format(dname, score_results[dname]))
    pickle.dump(score_results,
                open('../quality_evaluation_results/GPT3_score_{}.pkl'.format(gen_dir_name), 'wb+'))
    return score_results

def BERTScore(gen_dir_name):
    score_results = {}
    bert_score = evaluate.load('bertscore')
    model_type = input('Which BERT-based model will you use?')

    for root, ds, fs in os.walk("../model_output/{}".format(gen_dir_name)):
        for fn in fs:
            data = pd.read_csv(os.path.join(root, fn))
            dname = fn.split('.')[0]
            predictions = data['generated_text'].to_list()
            references = data['goldens'].to_list()

            f_n = []
            for p in predictions:
                f_n.append(str(p))

            result = bert_score.compute(predictions=f_n, references=references, model_type=model_type, batch_size=16)
            overall = sum(result['f1'])/len(result['f1'])
            print('Results for {} dataset: {}'.format(dname, overall))
            score_results[dname] = overall

    pickle.dump(score_results,
                open('../quality_evaluation_results/bert_score_{}.pkl'.format(gen_dir_name), 'wb+'))
    return score_results

def BARTscore(gen_dir_name, device):
    bart_scorer = BARTScorer(device=device, checkpoint='facebook/bart-large-cnn')
    bart_scorer.load(path='bart_score.pth')
    score_results = {}

    for root, ds, fs in os.walk("../model_output/{}".format(gen_dir_name)):
        for fn in fs:
            data = pd.read_csv(os.path.join(root, fn))
            dname = fn.split('.')[0]
            predictions = data['generated_text'].to_list()
            references = data['goldens'].to_list()

            f_n = []
            for p in predictions:
                f_n.append(str(p))

            result = bart_scorer.score(f_n, references)
            overall = sum(result) / len(result)
            print('Results for {} dataset: {}'.format(dname, overall))
            score_results[dname] = overall

    pickle.dump(score_results,
                open('../quality_evaluation_results/bart_score_{}.pkl'.format(gen_dir_name), 'wb+'))
    return score_results

def main(gen_dir_name: str, cuda: bool, device: str, score_method: str):
    if not os.path.exists("../quality_evaluation_results/"):
        os.mkdir("../quality_evaluation_results/")

    if score_method == 'bart_score':
        BARTscore(gen_dir_name, device)
    elif score_method == 'bert_score':
        BERTScore(gen_dir_name)
    elif score_method == 'GPT3_score':
        GPTScore(gen_dir_name)
    elif score_method == 'bleu':
        bleu(gen_dir_name)
    elif score_method == 'rouge':
        rouge(gen_dir_name)


if __name__ == '__main__':
    #IMHI_BARTscore('GPT4_expert')

    parser = argparse.ArgumentParser(
        description='The BART-score evaluation.')
    parser.add_argument('--gen_dir_name', type=str)
    parser.add_argument('--score_method', type=str, default='bart_score', choices=['bart_score', 'GPT3_score', 'bert_score', 'bleu', 'rouge'])
    parser.add_argument('--cuda', action='store_true')

    args = parser.parse_args()
    args = vars(args)
    device = torch.device("cuda:0" if torch.cuda.is_available() and args['cuda'] is True else "cpu")
    args['device'] = device

    main(**args)
src/score.py
SteveKGYang/MentalLLaMA