Multi-modal-Valuation-Forecast-System/datasets.py at main · Lanping-Tech/Multi-modal-Valuation-Forecast-System · GitHub

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
import torch
import torch.utils.data as data
import numpy as np
import os
from data_preprocessing import load_data
from transformers import AutoTokenizer


class MultiModalDataset(data.Dataset):
    def __init__(self, mts_data, text_data, window_size=5):
        self.window_size = window_size
        self.mts_data = mts_data
        self.text_data = text_data
        self.tokenizer = AutoTokenizer.from_pretrained('chinese-roberta-wwm-ext')
        # print(len(self.mts_data), len(self.text_data))

    def __getitem__(self, index):
        ts, text = self.mts_data[index], self.text_data[index]
        mts = ts[:self.window_size, :]
        mts = torch.tensor(mts, dtype=torch.float)
        label = ts[self.window_size, 3:4]
        label = torch.tensor(label, dtype=torch.float)
        content = []
        # print(len(text))
        for i in range(self.window_size):
            candidated_text = text[i]
            seleted_id = np.random.randint(len(candidated_text))
            content.append(candidated_text[seleted_id])

        input_ids_list = torch.empty(0, 128, dtype=torch.long)
        attention_mask_list = torch.empty(0, 128, dtype=torch.long)
        token_type_ids_list = torch.empty(0, 128, dtype=torch.long)

        for i in range(len(content)):
            content_encoding = self.tokenizer(content[i], add_special_tokens=True, max_length=128, padding='max_length', return_tensors='pt')
            input_ids = content_encoding['input_ids']
            attention_mask = content_encoding['attention_mask']
            token_type_ids = content_encoding['token_type_ids']
            input_ids_list = torch.cat((input_ids_list, input_ids), dim=0)
            attention_mask_list = torch.cat((attention_mask_list, attention_mask), dim=0)
            token_type_ids_list = torch.cat((token_type_ids_list, token_type_ids), dim=0)
        return mts, (input_ids_list, attention_mask_list, token_type_ids_list), label

    def __len__(self):
        return len(self.mts_data)


class MultiModalDataset_plus(data.Dataset):
    def __init__(self, mts_path, text_path, stock_ids, window_size=5):
        self.window_size = window_size
        self.mts_data, self.text_data, _, _ = load_data(mts_path, text_path, stock_ids, WINDOW_SIZE=window_size+1)
        self.tokenizer = AutoTokenizer.from_pretrained('chinese-roberta-wwm-ext')
        # print(len(self.mts_data), len(self.text_data))

    def __getitem__(self, index):
        ts, text = self.mts_data[index], self.text_data[index]
        mts = ts[:self.window_size, :]
        mts = torch.tensor(mts, dtype=torch.float)
        label = ts[self.window_size, 3:4]
        label = torch.tensor(label, dtype=torch.float)
        content = []
        # print(len(text))
        for i in range(self.window_size):
            candidated_text = text[i]
            seleted_id = np.random.randint(len(candidated_text))
            content.append(candidated_text[seleted_id])

        input_ids_list = torch.empty(0, 128, dtype=torch.long)
        attention_mask_list = torch.empty(0, 128, dtype=torch.long)
        token_type_ids_list = torch.empty(0, 128, dtype=torch.long)

        for i in range(len(content)):
            content_encoding = self.tokenizer(content[i], add_special_tokens=True, max_length=128, padding='max_length', return_tensors='pt')
            input_ids = content_encoding['input_ids']
            attention_mask = content_encoding['attention_mask']
            token_type_ids = content_encoding['token_type_ids']
            input_ids_list = torch.cat((input_ids_list, input_ids), dim=0)
            attention_mask_list = torch.cat((attention_mask_list, attention_mask), dim=0)
            token_type_ids_list = torch.cat((token_type_ids_list, token_type_ids), dim=0)
        return mts, (input_ids_list, attention_mask_list, token_type_ids_list), label

    def __len__(self):
        return len(self.mts_data)


if __name__ == '__main__':
    data = MultiModalDataset('data/股价', 'data/文本数据', ['000001'])
    for i in range(64):
        print(i, data[i][0].shape)
        print()