美股文本信息抽取

import re
import json


from rich import print
from transformers import AutoTokenizer, AutoModel

# 定义不同实体下的具备属性
schema = {
    '金融': ['日期', '股票名称', '开盘价', '收盘价', '成交量'],
}

# 信息抽取的模版
IE_PATTERN = "{}\n\n提取上述句子中{}的实体，并按照JSON格式输出，上述句子中不存在的信息用['原文中未提及']来表示，多个值之间用','分隔。"


# 提供一些例子供模型参考
ie_examples = {
        '金融': [
                    {
                        'content': '2023-01-10，股市震荡。股票古哥-D[EOOE]美股今日开盘价100美元，一度飙升至105美元，随后回落至98美元，最终以102美元收盘，成交量达到520000。',
                        'answers': {
                                        '日期': ['2023-01-10'],
                                        '股票名称': ['古哥-D[EOOE]美股'],
                                        '开盘价': ['100美元'],
                                        '收盘价': ['102美元'],
                                        '成交量': ['520000'],
                            }
                    }
        ]
}

# 定义init_prompts函数
def init_prompts():
    """
     初始化前置prompt，便于模型做 incontext learning。
     """
    ie_pre_history = [
        (
            "现在你需要帮助我完成信息抽取任务，当我给你一个句子时，你需要帮我抽取出句子中实体信息，并按照JSON的格式输出，上述句子中没有的信息用['原文中未提及']来表示，多个值之间用','分隔。",
            '好的，请输入您的句子。'
        )
    ]
    for _type, example_list in ie_examples.items():
        # print(f'_type-->{_type}')
        # print(f'example_list-->{example_list}')
        # print(f'*'*80)
        for example in example_list:
            sentence = example["content"]
            properties_str = ', '.join(schema[_type])
            # print(f'properties_str-->{properties_str}')
            schema_str_list = f'"{_type}"({properties_str})'
            # print(f'schema_str_list-->{schema_str_list}')

            sentence_with_prompt = IE_PATTERN.format(sentence, schema_str_list)
            # print(f'sentence_with_prompt-->{sentence_with_prompt}')
            ie_pre_history.append((f"{sentence_with_prompt}",f"{json.dumps(example['answers'], ensure_ascii=False)}"))
            # print(f'ie_pre_history-->{ie_pre_history}')

    return {"ie_pre_history":ie_pre_history}

def clean_response(response: str):
    """
    后处理模型输出。

    Args:
        response (str): _description_
    """
    if '```json' in response:
        res = re.findall(r'```json(.*?)```', response)
        if len(res) and res[0]:
            response = res[0]
        response = response.replace('、', ',')
    try:
        return json.loads(response)
    except:
        return response
        

def inference(sentences: list,
              custom_settings: dict):
    """
    推理函数。

    Args:
        sentences (List[str]): 待抽取的句子。
        custom_settings (dict): 初始设定，包含人为给定的 few-shot example。
    """
    for sentence in sentences:
        cls_res = "金融"
        if cls_res not in schema:
            print(f'The type model inferenced {cls_res} which is not in schema dict, exited.')
            exit()
        properties_str = ', '.join(schema[cls_res])
        schema_str_list = f'"{cls_res}"({properties_str})'
        sentence_with_ie_prompt = IE_PATTERN.format(sentence, schema_str_list)
        # print(f'sentence_with_prompt-->{sentence_with_ie_prompt}')
        ie_res, history = model.chat(tokenizer,
                                     sentence_with_ie_prompt,
                                     history=custom_settings["ie_pre_history"])
        ie_res = clean_response(ie_res)
        print(f'>>> [bold bright_red]sentence: {sentence}')
        print(f'>>> [bold bright_green]inference answer:{ie_res} ')



if __name__ == '__main__':
    #device = 'cuda:0'
    device = 'cpu'
    tokenizer = AutoTokenizer.from_pretrained("/Users/ligang/PycharmProjects/llm/ChatGLM-6B/THUDM/chatglm-6b-int4",
                                              trust_remote_code=True)
    #model = AutoModel.from_pretrained("./ChatGLM-6B/THUDM/chatglm-6b",
                                      # trust_remote_code=True).half().cuda()
    model = AutoModel.from_pretrained("/Users/ligang/PycharmProjects/llm/ChatGLM-6B/THUDM/chatglm-6b-int4",
                                      trust_remote_code=True).float()
    model.to(device)

    sentences = [
        '2023-02-15，寓意吉祥的节日，股票佰笃[BD]美股开盘价10美元，虽然经历了波动，但最终以13美元收盘，成交量微幅增加至460,000，投资者情绪较为平稳。',
        '2023-04-05，市场迎来轻松氛围，股票盘古(0021)开盘价23元，尽管经历了波动，但最终以26美元收盘，成交量缩小至310,000，投资者保持观望态度。',
    ]

    custom_settings = init_prompts()

    inference(
        sentences,
        custom_settings
    )



1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131