【文獻應用分享】FinGPT: Open-Source Financial Large Language Models

from transformers import AutoModel, AutoTokenizer
from peft import PeftModel
from torch import cuda
from sqlalchemy import create_engine
import pandas as pd

import time


def get_db_engine(user, password, ip, port, database):
    return create_engine(
        f'postgresql://{user}:{password}@{ip}:{port}/{database}')
    

def load_news_dataset(start_date, date_col, table_name, db_cofig):
    sql = f"""SELECT * FROM "{table_name}"
    WHERE "{date_col}" >= '{start_date}';"""
    
    return pd.read_sql(
        sql,
        get_db_engine(**db_cofig)
    )


def generate_prompts(titles):
    titles = pd.Series(titles)
    instruction_prompt = 'Instruction: What is the sentiment of this news? Please choose an answer from {negative/neutral/positive}'
    input_prompt = 'Input: '
    answer_prompt = 'Answer: '
    titles = titles.apply(lambda x: f'{instruction_prompt}\n{input_prompt}{x}\n{answer_prompt}')
    return titles.tolist()


def load_model(base_model_path: str, peft_model_path: str, cache_dir: str):
    tokenizer = AutoTokenizer.from_pretrained(
        base_model_path,
        trust_remote_code=True,
        cache_dir=cache_dir
    )
    model = AutoModel.from_pretrained(
        base_model_path,
        trust_remote_code=True, 
        device_map = "auto",
        cache_dir=cache_dir
    )
    model = PeftModel.from_pretrained(
        model,
        peft_model_path,
        cache_dir=cache_dir
    )
    model = model.eval()
    
    return tokenizer, model
    


def predict(prompts: list, model, tokenizer, device):
    print('---[Tokenize]...')
    start_time = time.time()
    tokens = tokenizer(prompts, return_tensors='pt', padding=True)
    print(f'------>[Tokenized] {time.time() - start_time}(s)')
    
    tokens.to(device)
    model = model.to(device)
    
    print('---[Generate Answer]...')
    res = model.generate(**tokens, max_length=512)
    print(f'------>[Generated] {time.time() - start_time}(s)')
    
    print('---[Decode and Get Key Answer]...')
    result_texts = []
    count = 0
    for i in res:
        print(f'------{count + 1}/{len(res)}')
        sentence = tokenizer.decode(i)
        key_answer = sentence.split('Answer: ')[1].replace(' ', '')
        result_texts.append(key_answer)
        count += 1
        
    # res_sentences = [tokenizer.decode(i) for i in res]
    # result_texts = [o.split("Answer: ")[1] for o in res_sentences]
    print(f'------>[Get] {time.time() - start_time}(s)')
    
    return result_texts


def process_batch(df, batch_prompts, model, tokenizer, device, db_table, db_cofig):
    # Predict sentiment.
    print(f'---[Predict Batch]...')
    start_time = time.time()
    result_texts = predict(batch_prompts, model, tokenizer, device)
    print(f'---> [Predicted Batch] {time.time() - start_time}(s)')
    # print(f'------> {result_texts}')
    
    # Save to DB.
    df['sentiment'] = result_texts
    save_to_db(df, db_table, db_cofig)


def save_to_db(df, db_table, db_cofig):
    df.to_sql(
        name=db_table,
        con=get_db_engine(**db_cofig),
        if_exists='append',
        index=False,
        chunksize=1000,
    )


if __name__ == '__main__':
    start_date = '2022-01-01'
    base_model_path = 'THUDM--chatglm2-6b'
    peft_model_path = 'oliverwang15--FinGPT_v31_ChatGLM2_Sentiment_Instruction_LoRA_FT/'
    cache_dir = './'
    load_db_table = ''
    save_db_table = ''
    db_date_col = 'date'
    db_title_col = 'headline'
    db_cofig = {
        'user': '', 
        'password': '',
        'ip': '',
        'port': '',
        'database': ''
    }
    
    # Get device.
    device = 'cuda' if cuda.is_available() else 'cpu'
    cuda.empty_cache()
    print(f'Devicd: {device}')

    # Get tokenizer, model.
    print('[Load Model and Tokenizer]...')
    start_time = time.time()
    tokenizer, model = load_model(
        base_model_path,
        peft_model_path,
        cache_dir
    )
    print(f'---> [Loaded] {time.time() - start_time}(s)')
    
    # Get dataset and generate promtps.
    print('[Load Dataset and Generate Prompts]...')
    start_time = time.time()
    
    df = load_news_dataset(start_date, db_date_col, load_db_table, db_cofig)
    prompts = generate_prompts(df[db_title_col])
    
    print(f'---> [Loaded and Generated] {time.time() - start_time}(s), {len(prompts)} rows.')
    
    # Divide to batch to predict.
    print(f'[Predict]...')
    start_time = time.time()
    batch_size = 100
    count = 0
    batch_num = round(len(prompts) / batch_size, 2)
    for i in range(0, len(prompts), batch_size):
        print(f'---(Batch {count + 1}/{batch_num})')
        batch = prompts[i:i + batch_size]
        process_batch(
            df.iloc[i:i + batch_size, :].copy(),
            batch,
            model,
            tokenizer,
            device,
            save_db_table,
            db_cofig
        )
        cuda.empty_cache()
        count += 1
    print(f'---> [Predicted] {time.time() - start_time}(s)')