forked from ProsusAI/finBERT
-
Notifications
You must be signed in to change notification settings - Fork 0
/
datasets.py
20 lines (15 loc) · 733 Bytes
/
datasets.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import argparse
import pandas as pd
from sklearn.model_selection import train_test_split
import os
if not os.path.exists('data/sentiment_data'):
os.makedirs('data/sentiment_data')
parser = argparse.ArgumentParser(description='Sentiment analyzer')
parser.add_argument('--data_path', type=str, help='Path to the text file.')
args = parser.parse_args()
data = pd.read_csv(args.data_path, sep='.@', names=['text','label'])
train, test = train_test_split(data, test_size=0.2, random_state=0)
train, valid = train_test_split(train, test_size=0.1, random_state=0)
train.to_csv('data/sentiment_data/train.csv',sep='\t')
test.to_csv('data/sentiment_data/test.csv',sep='\t')
valid.to_csv('data/sentiment_data/validation.csv',sep='\t')